算法学习|前缀和与差分

前缀和

快速计算数组的区间和，适用于多次区间查询的场景
时间复杂度：预处理 \(O(n)\)，查询 \(O(1)\)

def build_prefix_sum(arr):
    n = len(arr)
    prefix = [0] * (n + 1)
    for i in range(n):
        prefix[i+1] = prefix[i] + arr[i]
    return prefix

def query_prefix_sum(prefix, l, r):
    # 查询原数组区间 [l, r] 的和（索引从0开始）
    return prefix[r+1] - prefix[l]

arr = [3, 1, 4, 2, 5]
prefix = build_prefix_sum(arr)
print(prefix)  # 输出 [0, 3, 4, 8, 10, 15]

print(query_prefix_sum(prefix, 1, 3))

2维情况

def build_2d_prefix(matrix):
    rows = len(matrix)
    cols = len(matrix[0]) if rows > 0 else 0
    prefix = [[0]*(cols+1) for _ in range(rows+1)]
    for i in range(rows):
        for j in range(cols):
            prefix[i+1][j+1] = matrix[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j]
    return prefix

def query_2d_prefix(prefix, x1, y1, x2, y2):
    # 查询左上角(x1,y1)到右下角(x2,y2)的子矩阵和
    return prefix[x2+1][y2+1] - prefix[x1][y2+1] - prefix[x2+1][y1] + prefix[x1][y1]
	#这里所有输入均需要减一

差分

高效处理数组的区间更新操作（如区间加减），适用于多次区间修改后查询最终结果的场景

对于原数组 \(arr\)，差分数组 \(diff\) 满足：
\(diff[i] = arr[i] - arr[i-1]（i >= 1），diff[0] = arr[0]\)
原数组可通过前缀和从差分数组恢复：
\(arr[i] = diff[0] + diff[1] + ... + diff[i]\)

对原数组的区间 \([l, r]\) 增加 \(val\) 时，只需操作差分数组的两个位置：
\(diff[l] += val\)
\(diff[r+1] -= val\)

class DifferenceArray:
    def __init__(self, arr):
        self.n = len(arr)
        self.diff = [0] * self.n
        self.diff[0] = arr[0]
        for i in range(1, self.n):
            self.diff[i] = arr[i] - arr[i-1]

    def range_update(self, l, r, val):
        self.diff[l] += val
        if r + 1 < self.n:
            self.diff[r + 1] -= val

    def restore(self):
        res = [self.diff[0]]
        for i in range(1, self.n):
            res.append(res[-1] + self.diff[i])
        return res

arr = [1, 3, 5, 2, 4]
da = DifferenceArray(arr)
da.range_update(1, 3, 2)  # 区间 [1,3] 增加2
print(da.restore())  # 输出 [1, 5, 7, 4, 4]