算法刷题系列(四十五):线段树——维护区间查询与区间加更新
算法刷题系列(四十五):线段树——维护区间查询与区间加更新
算法刷题系列(四十五):线段树——维护区间查询与区间加更新
树状数组适合单点更新和前缀聚合;当题目需要灵活的区间操作时,线段树更通用。本文实现一个支持区间加、区间求和的懒标记线段树,区间统一使用半开形式 [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 进行授权