451_前缀和与差分数组¶

class NumArray:
    def __init__(self, nums):
        """
        预处理：计算前缀和数组
        prefix[i] = nums[0] + nums[1] + ... + nums[i-1]
        （多一位占位，prefix[0] = 0 方便统一处理左边界）
        """
        n = len(nums)
        self.prefix = [0] * (n + 1)        # 多一位，prefix[0]=0

        for i in range(n):
            self.prefix[i + 1] = self.prefix[i] + nums[i]  # 累加

    def sum_range(self, left: int, right: int) -> int:
        """
        O(1) 查询区间和
        区间[left, right]的和 = prefix[right+1] - prefix[left]
        """
        return self.prefix[right + 1] - self.prefix[left]

# 测试
na = NumArray([-2, 0, 3, -5, 2, -1])
print(na.sum_range(0, 2))  # 1（-2+0+3）
print(na.sum_range(2, 5))  # -1（3-5+2-1）
print(na.sum_range(0, 5))  # -3（全部求和）

class NumMatrix:
    def __init__(self, matrix):
        """
        二维前缀和：prefix[i][j] = matrix中(0,0)到(i-1,j-1)的矩形和
        容斥原理：大矩形 - 左边 - 上边 + 左上角（避免重复减）
        """
        m, n = len(matrix), len(matrix[0])
        self.prefix = [[0] * (n + 1) for _ in range(m + 1)]

        for i in range(1, m + 1):
            for j in range(1, n + 1):
                self.prefix[i][j] = (
                    matrix[i-1][j-1]          # 当前格子
                    + self.prefix[i-1][j]     # 上方矩形
                    + self.prefix[i][j-1]     # 左方矩形
                    - self.prefix[i-1][j-1]   # 减去重复的左上角
                )

    def sum_region(self, r1, c1, r2, c2) -> int:
        """
        O(1) 查询矩形区域和
        目标矩形 = 大矩形 - 左边 - 上边 + 左上角
        """
        return (
            self.prefix[r2+1][c2+1]
            - self.prefix[r1][c2+1]
            - self.prefix[r2+1][c1]
            + self.prefix[r1][c1]
        )

# 测试
matrix = [
    [3, 0, 1, 4, 2],
    [5, 6, 3, 2, 1],
    [1, 2, 0, 1, 5],
    [4, 1, 0, 1, 7],
    [1, 0, 3, 0, 5]
]
nm = NumMatrix(matrix)
print(nm.sum_region(2, 1, 4, 3))  # 8（右下角4×4子矩阵的某区域）
print(nm.sum_region(1, 1, 2, 2))  # 11（6+3+2+0）

def corp_flight_bookings(bookings, n):
    """
    差分数组：不逐一更新每个位置，而是在端点打标记
    最后对差分数组求前缀和，还原真实值

    差分数组 diff[i] = nums[i] - nums[i-1]
    区间 [l, r] 加 val：diff[l] += val, diff[r+1] -= val
    """
    diff = [0] * (n + 2)                   # 多两位防越界

    for first, last, seats in bookings:
        diff[first] += seats               # 区间左端+seats
        diff[last + 1] -= seats            # 区间右端+1 处-seats

    # 对差分数组求前缀和，还原真实预订量
    result = []
    running_sum = 0
    for i in range(1, n + 1):
        running_sum += diff[i]             # 累加得到第i个航班的总量
        result.append(running_sum)

    return result

# 测试
print(corp_flight_bookings([[1,2,10],[2,3,20],[2,5,25]], 5))
# [10, 55, 45, 25, 25]
# 航班1: 10
# 航班2: 10+20+25=55
# 航班3: 20+25=45
# 航班4: 25
# 航班5: 25

def subarray_sum(nums, k):
    """
    前缀和 + 哈希表
    子数组[i,j]的和 = prefix[j+1] - prefix[i]
    找 prefix[j+1] - prefix[i] = k，即找 prefix[j+1] - k 在之前出现过几次
    """
    count = 0
    prefix_sum = 0
    seen = {0: 1}                          # 前缀和→出现次数，初始化0出现1次

    for num in nums:
        prefix_sum += num                  # 更新当前前缀和

        # 找有多少个之前的前缀和等于 prefix_sum - k
        count += seen.get(prefix_sum - k, 0)

        seen[prefix_sum] = seen.get(prefix_sum, 0) + 1  # 记录当前前缀和

    return count

print(subarray_sum([1,1,1], 2))   # 2（[1,1]出现两次）
print(subarray_sum([1,2,3], 3))   # 2（[1,2]和[3]）

# 一维前缀和
n = len(nums)
prefix = [0] * (n + 1)
for i in range(n):
    prefix[i+1] = prefix[i] + nums[i]
# 查询 [l, r] 的和
range_sum = prefix[r+1] - prefix[l]

# 差分数组
diff = [0] * (n + 1)
# 区间 [l, r] 加 val
diff[l] += val
diff[r+1] -= val
# 还原
result = []
s = 0
for d in diff[:n]:
    s += d
    result.append(s)

# 前缀和 + 哈希（找和为k的子数组数量）
seen = {0: 1}
prefix_sum = count = 0
for num in nums:
    prefix_sum += num
    count += seen.get(prefix_sum - k, 0)
    seen[prefix_sum] = seen.get(prefix_sum, 0) + 1