文章

算法刷题系列(四十五):线段树——维护区间查询与区间加更新

算法刷题系列(四十五):线段树——维护区间查询与区间加更新

算法刷题系列(四十五):线段树——维护区间查询与区间加更新

树状数组适合单点更新和前缀聚合;当题目需要灵活的区间操作时,线段树更通用。本文实现一个支持区间加、区间求和的懒标记线段树,区间统一使用半开形式 [left, right)。

线段树的每个节点负责数组的一段区间。本文的合并函数是:

1
merge(leftSum, rightSum) = leftSum + rightSum

节点保存其负责区间的总和。若一个节点覆盖的整段都要增加 delta,其总和就增加 delta * 区间长度。懒标记记录尚未下传给子节点的增量,避免把更新拆到每个叶子。

查询和更新的不变量

递归处理节点区间 [lo, hi):

  • 与目标区间不相交:查询贡献为 0,更新不做事;
  • 完全包含于目标区间:查询直接返回节点和,更新直接改节点并记录懒标记;
  • 部分相交:先把懒标记下推,再递归左右子树,最后用合并函数重算父节点。

不变量:节点和始终等于该节点覆盖范围的真实元素总和;懒标记表示对子树整体生效、尚未逐层应用的统一增量。查询遇到部分重叠时先下推,才能安全访问子节点。

Java 实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
public class Main {
    static class SegmentTree {
        private final int n;
        private final long[] sum;
        private final long[] lazy;

        SegmentTree(long[] values) {
            if (values == null) throw new IllegalArgumentException("values must not be null");
            n = values.length;
            sum = new long[Math.max(1, 4 * n)];
            lazy = new long[Math.max(1, 4 * n)];
            if (n > 0) build(1, 0, n, values);
        }
        private void build(int node, int lo, int hi, long[] values) {
            if (hi - lo == 1) {
                sum[node] = values[lo];
                return;
            }
            int mid = lo + (hi - lo) / 2;
            build(node * 2, lo, mid, values);
            build(node * 2 + 1, mid, hi, values);
            pull(node);
        }
        private void pull(int node) {
            sum[node] = sum[node * 2] + sum[node * 2 + 1];
        }
        private void apply(int node, int length, long delta) {
            sum[node] += delta * length;
            lazy[node] += delta;
        }
        private void push(int node, int lo, int hi) {
            if (lazy[node] == 0 || hi - lo == 1) return;
            int mid = lo + (hi - lo) / 2;
            apply(node * 2, mid - lo, lazy[node]);
            apply(node * 2 + 1, hi - mid, lazy[node]);
            lazy[node] = 0;
        }
        void add(int left, int right, long delta) {
            checkRange(left, right);
            if (left < right) add(1, 0, n, left, right, delta);
        }
        private void add(int node, int lo, int hi, int ql, int qr, long delta) {
            if (ql <= lo && hi <= qr) {
                apply(node, hi - lo, delta);
                return;
            }
            push(node, lo, hi);
            int mid = lo + (hi - lo) / 2;
            if (ql < mid) add(node * 2, lo, mid, ql, qr, delta);
            if (qr > mid) add(node * 2 + 1, mid, hi, ql, qr, delta);
            pull(node);
        }
        long sum(int left, int right) {
            checkRange(left, right);
            return left == right ? 0 : sum(1, 0, n, left, right);
        }
        private long sum(int node, int lo, int hi, int ql, int qr) {
            if (ql <= lo && hi <= qr) return sum[node];
            push(node, lo, hi);
            int mid = lo + (hi - lo) / 2;
            long result = 0;
            if (ql < mid) result += sum(node * 2, lo, mid, ql, qr);
            if (qr > mid) result += sum(node * 2 + 1, mid, hi, ql, qr);
            return result;
        }
        private void checkRange(int left, int right) {
            if (left < 0 || left > right || right > n) {
                throw new IndexOutOfBoundsException("expected 0 <= left <= right <= n");
            }
        }
    }

    public static void main(String[] args) {
        SegmentTree tree = new SegmentTree(new long[]{2, 1, 3, 4});
        System.out.println(tree.sum(1, 4)); // 8
        tree.add(1, 3, 5);                 // 数组变为 [2, 6, 8, 4]
        System.out.println(tree.sum(0, 4)); // 20
        System.out.println(tree.sum(2, 2)); // 0
    }
}

Python 实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class SegmentTree:
    def __init__(self, values: list[int]):
        if values is None:
            raise ValueError("values must not be None")
        self.n = len(values)
        size = max(1, 4 * self.n)
        self.total = [0] * size
        self.lazy = [0] * size
        if self.n:
            self._build(1, 0, self.n, values)

    def _build(self, node: int, lo: int, hi: int, values: list[int]) -> None:
        if hi - lo == 1:
            self.total[node] = values[lo]
            return
        mid = (lo + hi) // 2
        self._build(node * 2, lo, mid, values)
        self._build(node * 2 + 1, mid, hi, values)
        self._pull(node)

    def _pull(self, node: int) -> None:
        self.total[node] = self.total[node * 2] + self.total[node * 2 + 1]

    def _apply(self, node: int, length: int, delta: int) -> None:
        self.total[node] += delta * length
        self.lazy[node] += delta

    def _push(self, node: int, lo: int, hi: int) -> None:
        if self.lazy[node] == 0 or hi - lo == 1:
            return
        mid = (lo + hi) // 2
        self._apply(node * 2, mid - lo, self.lazy[node])
        self._apply(node * 2 + 1, hi - mid, self.lazy[node])
        self.lazy[node] = 0

    def _check(self, left: int, right: int) -> None:
        if not 0 <= left <= right <= self.n:
            raise IndexError("expected 0 <= left <= right <= n")

    def add(self, left: int, right: int, delta: int) -> None:
        self._check(left, right)
        if left < right:
            self._add(1, 0, self.n, left, right, delta)

    def _add(self, node: int, lo: int, hi: int, ql: int, qr: int, delta: int) -> None:
        if ql <= lo and hi <= qr:
            self._apply(node, hi - lo, delta)
            return
        self._push(node, lo, hi)
        mid = (lo + hi) // 2
        if ql < mid:
            self._add(node * 2, lo, mid, ql, qr, delta)
        if qr > mid:
            self._add(node * 2 + 1, mid, hi, ql, qr, delta)
        self._pull(node)

    def sum(self, left: int, right: int) -> int:
        self._check(left, right)
        if left == right:
            return 0
        return self._sum(1, 0, self.n, left, right)

    def _sum(self, node: int, lo: int, hi: int, ql: int, qr: int) -> int:
        if ql <= lo and hi <= qr:
            return self.total[node]
        self._push(node, lo, hi)
        mid = (lo + hi) // 2
        result = 0
        if ql < mid:
            result += self._sum(node * 2, lo, mid, ql, qr)
        if qr > mid:
            result += self._sum(node * 2 + 1, mid, hi, ql, qr)
        return result

tree = SegmentTree([2, 1, 3, 4])
assert tree.sum(1, 4) == 8
tree.add(1, 3, 5)
assert tree.sum(0, 4) == 20
assert tree.sum(2, 2) == 0
empty = SegmentTree([])
assert empty.sum(0, 0) == 0

复杂度、边界与扩展

建树 O(n);区间加和区间和均为 O(log n)(覆盖的规范区间节点数为对数级),空间 O(n)。空数组允许构造,但只有空区间 [0,0) 合法。区间采用左闭右开,left == right 是合法空区间,求和为 0、更新无影响。

本文的 merge 是求和;如果改为最大值,pull 应取左右最大,查询的空区间单位元应为负无穷,不能仍返回 0。懒标记的 apply 也要匹配聚合语义。常见错误包括遗漏区间长度因子、部分覆盖前不下推、递归后不重算父节点、把左右边界当闭区间,以及对空区间递归导致无限拆分。long 求和仍需确认题目范围不会溢出。

小结

线段树的两个关键设计点是节点维护什么聚合,以及父节点如何 merge 子节点。懒传播则利用“整段统一更新可直接更新聚合值”的性质,把区间操作从逐元素更新缩短为对数级。

本文由作者按照 CC BY 4.0 进行授权