Solution
from typing import List
class BinaryIndexedTree:
__slots__ = ["n", "c"]
def __init__(self, n):
self.n = n
self.c = [0] * (n + 1)
def update(self, x: int, delta: int):
while x <= self.n:
self.c[x] += delta
x += x & -x
def query(self, x: int) -> int:
s = 0
while x > 0:
s += self.c[x]
x -= x & -x
return s
class NumArray:
__slots__ = ["tree"]
def __init__(self, nums: List[int]):
self.tree = BinaryIndexedTree(len(nums))
for i, v in enumerate(nums, 1):
self.tree.update(i, v)
def update(self, index: int, val: int) -> None:
prev = self.sumRange(index, index)
self.tree.update(index + 1, val - prev)
def sumRange(self, left: int, right: int) -> int:
return self.tree.query(right + 1) - self.tree.query(left)
if __name__ == "__main__":
arr = NumArray([1, 3, 5])
print(arr.sumRange(0, 2)) #9
arr.update(1, 2)
print(arr.sumRange(0, 2)) #8
Solution 2
from typing import List
class Node:
__slots__ = ["l", "r", "v"]
def __init__(self):
self.l = self.r = self.v = 0
class SegmentTree:
__slots__ = ["nums", "tr"]
def __init__(self, nums):
self.nums = nums
n = len(nums)
self.tr = [Node() for _ in range(n << 2)]
self.build(1, 1, n)
def build(self, u, l, r):
self.tr[u].l, self.tr[u].r = l, r
if l == r:
self.tr[u].v = self.nums[l - 1]
return
mid = (l + r) >> 1
self.build(u << 1, l, mid)
self.build(u << 1 | 1, mid + 1, r)
self.pushup(u)
def modify(self, u, x, v):
if self.tr[u].l == x and self.tr[u].r == x:
self.tr[u].v = v
return
mid = (self.tr[u].l + self.tr[u].r) >> 1
if x <= mid:
self.modify(u << 1, x, v)
else:
self.modify(u << 1 | 1, x, v)
self.pushup(u)
def query(self, u, l, r):
if self.tr[u].l >= l and self.tr[u].r <= r:
return self.tr[u].v
mid = (self.tr[u].l + self.tr[u].r) >> 1
result = 0
if l <= mid:
result += self.query(u << 1, l, r)
if r > mid:
result += self.query(u << 1 | 1, l, r)
return result
def pushup(self, u):
self.tr[u].v = self.tr[u << 1].v + self.tr[u << 1 | 1].v
class NumArray:
__slots__ = ["tree"]
def __init__(self, nums: List[int]):
self.tree = SegmentTree(nums)
def update(self, index: int, val: int) -> None:
self.tree.modify(1, index + 1, val)
def sumRange(self, left: int, right: int) -> int:
return self.tree.query(1, left + 1, right + 1)
if __name__ == "__main__":
arr = NumArray([1, 3, 5])
print(arr.sumRange(0, 2))
arr.update(1, 2)
print(arr.sumRange(0, 2))