307. 区域和检索 - 数组可修改-EXIAOHU Blog

题目链接：307. 区域和检索 - 数组可修改

线段树。使用一棵二叉树，这棵二叉树中结点的值为数组中某一个范围内的和。

树的根结点保存数组的总和 sum(arr, 0, n -1) ；
将父结点对应的数列 [s, e] 二分为 [s, m] 和 [m + 1, e]：
- 其左子树的根结点保存 sum(arr, s, m)；
- 右子树的根结点保存 sum(arr, m + 1, e)；

这样，update 和 sumRange 操作都可以在 $log(n)$ 的时间复杂度内完成。

from typing import List


class NumArray:
    def __init__(self, nums: List[int]):
        self.n = len(nums)
        self._seg = [0] * (len(nums) * 4)
        self._build(nums, 0, 0, len(nums) - 1)

    def _build(self, nums: List[int], node: int, s: int, e: int):
        if s == e:
            self._seg[node] = nums[s]
        else:
            m = s + (e - s) // 2
            self._build(nums, node * 2 + 1, s, m)
            self._build(nums, node * 2 + 2, m + 1, e)
            self._seg[node] = self._seg[node * 2 + 1] + self._seg[node * 2 + 2]

    def _change(self, index: int, val: int, node: int, s: int, e: int):
        if s == e:
            self._seg[node] = val
            return
        m = s + (e - s) // 2
        if index <= m:
            self._change(index, val, node * 2 + 1, s, m)
        else:
            self._change(index, val, node * 2 + 2, m + 1, e)
        self._seg[node] = self._seg[node * 2 + 1] + self._seg[node * 2 + 2]

    def _range(self, left: int, right: int, node: int, s: int, e: int) -> int:
        if left == s and right == e:
            return self._seg[node]
        m = s + (e - s) // 2
        if right <= m:
            return self._range(left, right, node * 2 + 1, s, m)
        if left > m:
            return self._range(left, right, node * 2 + 2, m + 1, e)
        return self._range(left, m, node * 2 + 1, s, m) + self._range(m + 1, right, node * 2 + 2, m + 1, e)

    def update(self, index: int, val: int) -> None:
        self._change(index, val, 0, 0, self.n - 1)

    def sumRange(self, left: int, right: int) -> int:
        return self._range(left, right, 0, 0, self.n - 1)