题目描述

给定某个数列的前缀和 $P$, 你需要支持以下两种操作：

1. 给定 $i$ 和新值 $v$，将 $P_i$ 改变为 $v$。
2. 给定 $l$ 和 $r$，求 $\sum_{i=l}^r P_i$。

函数签名

class PrefixSum:
    def __init__(self, arr: List[int])
    def update(self, i: int, v: int) -> None:
    def query(self, l: int, r: int) -> int:

输入

构造函数 PrefixSum 接收一个整数数组 arr，表示原始数列。

操作函数 update 接收两个整数 i 和 v，表示需要修改前缀和的下标和新值。

操作函数 query 接收两个整数 l 和 r，表示需要求和的区间左右端点。

输出

对于每个操作函数 query，返回区间 $[l, r]$ 的前缀和的总和。

示例

preSum = PrefixSum([1, 3, 4, 8, 6, 1])
preSum.query(2, 4) # 16
preSum.update(3, 2)
preSum.query(2, 4) # 14

思路

这道题目我们很容易想到可以使用前缀和来解决。首先我们在构造函数中对原始数列进行前缀和的预处理，使得我们能够在 $O(1)$ 的时间复杂度内计算任意区间的和。

对于 update 操作，我们只需要将下标为 $i$ 的位置改变为 $v$，而后对于大于等于 $i$ 的所有位置，都需要相应地更新其前缀和。

对于 query 操作，我们只需要用右下角的前缀和减去左上角不在目标区间内的前缀和，就可以获取到目标区间的总和。

代码实现

class PrefixSum:
    def __init__(self, arr: List[int]):
        self.preSum = [0]
        for num in arr:
            self.preSum.append(self.preSum[-1] + num)
    
    def update(self, i: int, v: int) -> None:
        diff = v - (self.preSum[i+1] - self.preSum[i])
        for j in range(i+1, len(self.preSum)):
            self.preSum[j] += diff
    
    def query(self, l: int, r: int) -> int:
        return self.preSum[r+1] - self.preSum[l]