线段树¶
线段树把一个大区间递归拆成左右两部分,每个节点保存所负责区间的统计结果。查询时,只合并与目标区间相关的节点;更新时,只修改从叶子到根的路径。
下面实现单点赋值与区间和查询。用求和作为例子,重点理解区间如何拆分和合并。
先统一接口约定¶
使用 0 开始的下标,查询采用左闭右开区间 [left, right)。例如查询前三项是 query(0, 3),空区间的和为 0。
数组表示的迭代实现中,叶子从位置 size 开始。size 是不小于 n 的最小 2 的幂,未使用的叶子填 0。
class SegmentTree:
def __init__(self, values):
self.n = len(values)
self.size = 1
while self.size < self.n:
self.size *= 2
self.tree = [0] * (2 * self.size)
for i, value in enumerate(values):
self.tree[self.size + i] = value
for i in range(self.size - 1, 0, -1):
self.tree[i] = self.tree[2 * i] + self.tree[2 * i + 1]
def set_value(self, index, value):
if not 0 <= index < self.n:
raise IndexError("update index out of range")
i = self.size + index
self.tree[i] = value
i //= 2
while i:
self.tree[i] = self.tree[2 * i] + self.tree[2 * i + 1]
i //= 2
def query(self, left, right):
if not 0 <= left <= right <= self.n:
raise IndexError("invalid query interval")
left += self.size
right += self.size
result = 0
while left < right:
if left % 2:
result += self.tree[left]
left += 1
if right % 2:
right -= 1
result += self.tree[right]
left //= 2
right //= 2
return result
seg = SegmentTree([2, 1, 3, 4, 5])
assert seg.query(1, 4) == 8
seg.set_value(2, 10)
assert seg.query(1, 4) == 15
assert seg.query(2, 2) == 0
assert SegmentTree([]).query(0, 0) == 0
查询为什么这样移动边界?¶
- left 是右子节点时,它的左兄弟不在查询范围内,所以先取当前节点,再向右移动。
- right 是右边界且不包含在范围内。它是右子节点时,先向左移动,再取对应节点。
- 处理完边缘节点后,把左右边界一起移到父层,继续合并更大的完整区间。
每一层至多处理两个边缘节点,树高为 O(log n),因此单次查询和单点更新耗时 O(log n)。建树和存储空间都是 O(n)。
合并规则决定能做什么¶
区间和用加法合并,空区间对应 0;区间最小值用 min 合并,空区间通常对应正无穷。可结合的合并规则,让相邻区间的答案能拼成更大区间的答案。
本例的查询写法适用于求和。若合并操作不满足交换律,例如字符串拼接,需要分别维护左右累积结果,并按原有顺序合并,不能直接套用单个 result。
区间更新是另一个层次
这份实现不包含区间批量更新。需要对一整段区间增加某个值时,可继续学习懒标记:延后下传更新,同时维护节点统计值。
与树状数组比较¶
| 需求 | 通常可以选择 |
|---|---|
| 静态区间和 | 前缀和 |
| 单点增加 + 动态区间和 | 树状数组 |
| 单点赋值 + 区间最值/和 | 线段树 |
| 区间批量更新 + 区间查询 | 支持对应懒标记的线段树 |
选择数据结构之前,先把更新和查询的类型写清楚。支持更多操作的结构,也意味着更多实现与维护成本。