高级数据结构:线段树与树状数组
以 (O(\log n)) 时间完成区间查询与修改,解决数组区间问题的高效利器。
线段树(Segment Tree)和树状数组(Binary Indexed Tree / Fenwick Tree)是处理区间查询和更新问题的核心数据结构。面试中常用于考查候选人对复杂数据结构的编码能力和复杂度分析。
树状数组(Binary Indexed Tree)
树状数组是线段树的轻量替代,适用于单点修改 + 前缀和查询的场景。
核心思想
利用二进制索引的性质,每个元素 tree[i] 存储从 i - lowbit(i) + 1 到 i 的区间和。
index: 1 2 3 4 5 6 7 8
│ │ │ │ │ │ │ │
1 2 3 4 5 6 7 8
└──┬─┘ └──┬─┘ └──┬─┘ └──┬─┘
2 3 6 7
└────┬────┘ └────┬────┘
4 6
└────────┬──────────┘
8
基础实现:单点更新 + 前缀和查询
class BIT:
"""树状数组:单点更新 + 前缀和查询"""
def __init__(self, n):
self.n = n
self.tree = [0] * (n + 1) # 1-indexed
def lowbit(self, x):
return x & (-x)
# 单点增加:在下标 idx 处增加 delta
def add(self, idx, delta):
while idx <= self.n:
self.tree[idx] += delta
idx += self.lowbit(idx)
# 前缀和查询:[1, idx]
def query(self, idx):
res = 0
while idx > 0:
res += self.tree[idx]
idx -= self.lowbit(idx)
return res
# 区间和查询:[l, r]
def range_query(self, l, r):
return self.query(r) - self.query(l - 1)
# 初始化:O(n log n)
def build(self, arr):
for i, val in enumerate(arr, 1):
self.add(i, val)
时间复杂度
| 操作 | 复杂度 | 原因 |
|---|---|---|
| 单点更新 | (O(\log n)) | 更新 lowbit 链上的节点 |
| 前缀和查询 | (O(\log n)) | 累加 lowbit 分解的区间 |
| 区间查询 | (O(\log n)) | 两次前缀和相减 |
| 建树 | (O(n \log n)) | 或 (O(n)) 线性构建 |
适用场景
| 问题类型 | BIT 解法 |
|---|---|
| 区间和查询 + 单点修改 | 标准 BIT |
| 逆序对计数 | 数值离散化 + BIT |
| 二维区间和 | 二维 BIT |
| 动态频率统计 | 实时频率前缀和 |
线段树(Segment Tree)
线段树是更通用的数据结构,支持区间修改 + 区间查询。
结构
- 每个节点代表一个区间
[l, r] - 叶子节点代表单个元素
- 内部节点存储该区间的聚合值(和、最值、计数等)
[0, 7] sum=36
/ \
[0, 3] sum=10 [4, 7] sum=26
/ \
[0,1] [2,3]
/ \
[0,0] [1,1]
基础实现:区间和查询 + 单点更新
class SegmentTree:
"""线段树:区间和查询 + 单点更新"""
def __init__(self, arr):
self.n = len(arr)
self.tree = [0] * (4 * self.n) # 4n 足够容纳完全二叉树
self._build(arr, 1, 0, self.n - 1)
def _build(self, arr, node, l, r):
if l == r:
self.tree[node] = arr[l]
return
mid = (l + r) // 2
self._build(arr, node * 2, l, mid)
self._build(arr, node * 2 + 1, mid + 1, r)
self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]
# 单点更新:将下标 idx 的值改为 val
def update(self, idx, val, node=1, l=0, r=None):
if r is None:
r = self.n - 1
if l == r:
self.tree[node] = val
return
mid = (l + r) // 2
if idx <= mid:
self.update(idx, val, node * 2, l, mid)
else:
self.update(idx, val, node * 2 + 1, mid + 1, r)
self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]
# 区间查询:查询 [ql, qr] 的和
def query(self, ql, qr, node=1, l=0, r=None):
if r is None:
r = self.n - 1
if ql <= l and r <= qr:
return self.tree[node]
if r < ql or l > qr:
return 0
mid = (l + r) // 2
left_sum = self.query(ql, qr, node * 2, l, mid)
right_sum = self.query(ql, qr, node * 2 + 1, mid + 1, r)
return left_sum + right_sum
Lazy Propagation(区间修改的加速)
当需要区间批量修改(如给区间 [l, r] 每个元素加 val)时,如果每次都遍历到叶子节点更新,复杂度为 (O(n \log n))。
Lazy 思想:延迟下传更新标记,只在必要时才将标记下推到子节点。
class LazySegmentTree:
"""带 Lazy Propagation 的线段树:区间修改 + 区间查询"""
def __init__(self, arr):
self.n = len(arr)
self.tree = [0] * (4 * self.n)
self.lazy = [0] * (4 * self.n) # 延迟标记
self._build(arr, 1, 0, self.n - 1)
def _build(self, arr, node, l, r):
if l == r:
self.tree[node] = arr[l]
return
mid = (l + r) // 2
self._build(arr, node * 2, l, mid)
self._build(arr, node * 2 + 1, mid + 1, r)
self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]
def _push(self, node, l, r):
"""下传 lazy 标记到子节点"""
if self.lazy[node] != 0:
mid = (l + r) // 2
# 更新左子树
self.tree[node * 2] += self.lazy[node] * (mid - l + 1)
self.lazy[node * 2] += self.lazy[node]
# 更新右子树
self.tree[node * 2 + 1] += self.lazy[node] * (r - mid)
self.lazy[node * 2 + 1] += self.lazy[node]
# 清空当前节点标记
self.lazy[node] = 0
# 区间修改:给 [ql, qr] 区间内每个元素增加 val
def range_add(self, ql, qr, val, node=1, l=0, r=None):
if r is None:
r = self.n - 1
if ql <= l and r <= qr:
self.tree[node] += val * (r - l + 1)
self.lazy[node] += val
return
if r < ql or l > qr:
return
self._push(node, l, r)
mid = (l + r) // 2
self.range_add(ql, qr, val, node * 2, l, mid)
self.range_add(ql, qr, val, node * 2 + 1, mid + 1, r)
self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]
# 区间查询:[ql, qr] 的和
def range_query(self, ql, qr, node=1, l=0, r=None):
if r is None:
r = self.n - 1
if ql <= l and r <= qr:
return self.tree[node]
if r < ql or l > qr:
return 0
self._push(node, l, r)
mid = (l + r) // 2
return (self.range_query(ql, qr, node * 2, l, mid) +
self.range_query(ql, qr, node * 2 + 1, mid + 1, r))
时间复杂度对比
| 操作 | 普通线段树 | Lazy 线段树 |
|---|---|---|
| 单点更新 | (O(\log n)) | (O(\log n)) |
| 区间查询 | (O(\log n)) | (O(\log n)) |
| 区间修改 | (O(n \log n)) | (O(\log n)) |
线段树 vs 树状数组
| 特性 | 树状数组(BIT) | 线段树 |
|---|---|---|
| 代码量 | 20 行左右 | 80+ 行 |
| 区间修改 | 需差分转化 | 原生支持(Lazy) |
| 区间查询 | 前缀和相减 | 原生支持 |
| 功能扩展 | 难(仅限可减可加) | 易(最值、计数、覆盖等) |
| 常数 | 极小 | 中等 |
| 面试手撕 | 推荐(简单清晰) | 展示实力 |
面试建议:
- 如果只是区间和 + 单点修改,写 BIT 即可
- 如果需要区间修改或复杂聚合,写 Lazy 线段树
- 面试时可以两种都提,优先实现 BIT
table
经典面试题
题目 1:Range Sum Query - Mutable
LeetCode 307:实现一个支持单点更新和区间和查询的数据结构。
BIT 解法:
class NumArray:
def __init__(self, nums):
self.n = len(nums)
self.nums = nums
self.bit = BIT(self.n)
for i, val in enumerate(nums, 1):
self.bit.add(i, val)
def update(self, index, val):
delta = val - self.nums[index]
self.nums[index] = val
self.bit.add(index + 1, delta)
def sumRange(self, left, right):
return self.bit.range_query(left + 1, right + 1)
题目 2:Count of Smaller Numbers After Self
LeetCode 315:统计数组中每个元素右边比它小的元素个数。
思路:离散化 + BIT,从右向左遍历,查询已出现元素中小于当前值的数量。
class Solution:
def countSmaller(self, nums):
# 离散化
sorted_unique = sorted(set(nums))
rank = {v: i + 1 for i, v in enumerate(sorted_unique)} # 1-indexed
bit = BIT(len(sorted_unique))
res = []
for num in reversed(nums):
r = rank[num]
# 查询 rank < r 的数量
res.append(bit.query(r - 1))
bit.add(r, 1)
return res[::-1]
题目 3:Range Module
LeetCode 715:动态添加/删除区间,查询某点是否被覆盖。
线段树解法:用 Lazy 线段树维护区间覆盖状态,支持区间赋值和单点查询。
面试追问
| 追问 | 回答要点 |
|---|---|
| “BIT 和线段树怎么选?” | BIT 代码短、常数小,适合简单区间和;线段树更通用,支持区间修改和复杂聚合 |
| “BIT 区间修改怎么做?” | 差分数组 + 两个 BIT 维护,或用树状数组维护差分的前缀和 |
| “线段树空间为什么开 4n?” | 完全二叉树最坏情况下需要约 4n 空间 |
| “Lazy 标记不下传会怎样?” | 查询结果错误,因为父节点的值已经更新但子节点还是旧值 |
| “逆序对除了 BIT 还能怎么做?” | 归并排序统计(O(n log n)),或平衡树 / Order Statistic Tree |
总结速查
| 场景 | 推荐方案 | 复杂度 |
|---|---|---|
| 区间和 + 单点修改 | BIT | 查询/修改 O(log n) |
| 区间和 + 区间修改 | Lazy 线段树 | 查询/修改 O(log n) |
| 区间最值 + 单点修改 | 线段树 | 查询/修改 O(log n) |
| 区间覆盖/线段染色 | Lazy 线段树 | 查询/修改 O(log n) |
| 逆序对计数 | BIT + 离散化 | O(n log n) |
| 二维区间和 | 二维 BIT / 二维线段树 | O(log² n) |
继续阅读
探索更多技术文章
浏览归档,发现更多关于系统设计、工具链和工程实践的内容。