跳转至

线段树

线段树把一个大区间递归拆成左右两部分,每个节点保存所负责区间的统计结果。查询时,只合并与目标区间相关的节点;更新时,只修改从叶子到根的路径。

下面实现单点赋值与区间和查询。用求和作为例子,重点理解区间如何拆分和合并。

先统一接口约定

使用 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。

区间更新是另一个层次

这份实现不包含区间批量更新。需要对一整段区间增加某个值时,可继续学习懒标记:延后下传更新,同时维护节点统计值。

与树状数组比较

需求 通常可以选择
静态区间和 前缀和
单点增加 + 动态区间和 树状数组
单点赋值 + 区间最值/和 线段树
区间批量更新 + 区间查询 支持对应懒标记的线段树

选择数据结构之前,先把更新和查询的类型写清楚。支持更多操作的结构,也意味着更多实现与维护成本。