LintCode 840. 可变范围求和
题目描述
给定一个整数数组 nums, 然后你需要实现两个函数：
-- update(i, val) 将数组下标为i的元素修改为val
-- sumRange(l, r) 返回数组下标在[l, r][l,r]区间的元素的和
注意事项
数组只能通过update函数进行修改。
你可以假设 update 函数与 sumRange 函数的调用数量是均匀的。
样例
输入:
nums = [1, 3, 5]
sumRange(0, 2)
update(1, 2)
sumRange(0, 2)
输出:
9
8
输入:
nums = [0, 9, 5, 7, 3]
sumRange(4, 4)
sumRange(2, 4)
update(4, 5)
update(1, 7)
update(0, 8)
sumRange(1, 2)
输出:
3
15
12
思路
列表切片返回区间和是行不通的，别天真了，会超时，这题可以用线段树来解。
树节点代表的是一个区间[left, right]，存储的是区间和，叶子节点存储的是数组中的每个数字。除叶子节点外，每个树节点都会划分2个子区间[left, mid], [mid+1, right]分配给其左、右儿子，其中mid就是left和right的中点，于是就有这个性质：每个节点的值是其左、右儿子的值之和。从而在更新节点值时，我们需要对应地将所有关联的父节点也更新。
代码
1.
class NumArray:
class TreeNode:
def __init__(self, start, end, val):
"""每个节点存储的是由start~end位置的数字之和，
左儿子是start~mid，右儿子是mid+1~end，
叶子节点的start=end"""
self.start = start
self.end = end
self.val = val
self.left = None
self.right = None
def __init__(self, nums):
"""
:type nums: List[int]
"""
self.root = self._build_tree(nums, 0, len(nums) - 1)
def update(self, i, val):
"""
:type i: int
:type val: int
:rtype: void
"""
self._modify(self.root, i, val)
def sumRange(self, i, j):
"""
:type i: int
:type j: int
:rtype: int
"""
return self._query(self.root, i, j)
def _build_tree(self, nums, start, end):
node = self.TreeNode(start, end, nums[start])
# 叶子节点
if start == end:
return node
mid = (start + end) // 2
node.left = self._build_tree(nums, start, mid)
node.right = self._build_tree(nums, mid + 1, end)
# 本节点存储的值等于其左儿子与右儿子存储的值之和
node.val = node.left.val + node.right.val
return node
def _modify(self, node, i, val):
"""修改位置i的节点值"""
# 已找到对应的叶子节点，直接更新节点值
if node.start == node.end:
node.val = val
return
if node.left.end < i:
self._modify(node.right, i, val)
else:
self._modify(node.left, i, val)
# 最后要更新本节点的值
node.val = node.left.val + node.right.val
def _query(self, node, i, j):
"""查询位置区间[i,j]的数字之和"""
# 若节点代表的区间被包含在查询区间内，则直接返回其存储的值
if i <= node.start and node.end <= j:
return node.val
# 节点代表的区间与查询区间无交集
if node.start > j or node.end < i:
return 0
# 节点代表的区间与查询区间有交集，则需要进一步在其左、右儿子中查询
return self._query(node.left, i, j) + self._query(node.right, i, j)
2.
class NumArray:
def __init__(self, nums):
"""
:type nums: List[int]
"""
self.size = len(nums)
# nums中的数字用作叶子节点，同时对应构造等数量的父节点，
# 用于存储区间和
self.tree = [0] * self.size + nums
for i in range(self.size - 1, 0, -1):
# 左儿子2i、右儿子2i+1
self.tree[i] = self.tree[i << 1] + self.tree[i << 1 | 1]
def update(self, i, val):
"""
:type i: int
:type val: int
:rtype: void
"""
i += self.size
# 先更新对应的叶子节点
self.tree[i] = val
while i > 1:
# 当i是偶数时，代表左儿子，i XOR 1则为i+1，即右儿子；
# 当i是奇数时，代表右儿子，i XOR 1 则为i-1，即左儿子
self.tree[i >> 1] = self.tree[i] + self.tree[i ^ 1]
i >>= 1
def sumRange(self, i, j):
"""
:type i: int
:type j: int
:rtype: int
"""
# 从叶子节点回溯
i += self.size
j += self.size
s = 0
while i <= j:
# i是奇数，代表右儿子
if i & 1:
s += self.tree[i]
# 变为后一个子区间的左儿子
i += 1
# 后一个子区间的父节点
i >>= 1
# j是偶数，代表左儿子
if not j & 1:
s += self.tree[j]
# 变为前一个子区间的右儿子
j -= 1
# 前一个子区间的父节点
j >>= 1
return s
