数据结构学习笔记 - 10 线段树
介绍
经典面试问题,一面墙,每次可以在其中一段区间进行染色,染色可以覆盖之前的染色,若干次染色后,问某个区间内有几种颜色?
这里涉及了两个操作,一个是更新,一个是查询。可以使用数组解决这个问题,染色就修改数组的指定区间,查询就是遍历整个数组,两种操作的复杂度都为 O(n)。
而如果使用线段树解决,复杂度为 O(logn)。
| 操作 | 使用数组复杂度 | 使用线段树复杂度 |
|---|---|---|
| 区间更新 | O(n) | O(logn) |
| 区间查询 | O(n) | O(logn) |
对于线段树,不考虑添加和删除操作,线段树解决的问题,区间本身是固定的,因此存储线段树使用静态数组即可。假定研究的问题是数组区间求和,数组共 8 个元素,则构造的线段树如下:
A[0...7]
/ \
A[0...3] A[4...7]
/ \ / \
A[0...1] A[2...3] A[4...5] A[6...7]
/ \ / \ / \ / \
A[0] A[1] A[2] A[3] A[4] A[5] A[6] A[7]
// 根节点存储全部区间的和,A[0...3]存储前半段的和,以此类推。
如果数组是 10 个元素,线段树则表示为:
A[0...9]
/ \
A[0...4] A[5...9]
/ \ / \
A[0...1] A[2...4] A[5...6] A[7...9]
/ \ / \ / \ / \
A[0] A[1] A[2] A[3,4] A[5] A[6] A[7] A[8,9]
/ \ / \
A[3] A[4] A[8] A[9]
- 线段树不是满二叉树,更不是完全二叉树。
- 平衡二叉树,指二叉树的每个节点的左右子树的高度差的绝对值不超过 1。平衡二叉树一定不会退化成链表。
- 堆就是平衡二叉树,完全二叉树就一定是平衡二叉树。但是二分搜索树就不一定是平衡二叉树。
- 线段树是一棵平衡二叉树。
- 线段树虽然不是完全二叉树,但类似也是只有最下层是不「满」的,我们可以近似将其看作一个满二叉树,因此也可以使用数组表示。
- 对于线段树来说,如果使用数组来表示,如果区间有 n 个元素,那么数组需要 4n 的空间来存储。对于线段树我们不考虑添加元素,即区间固定。所以创建线段树复杂度为 O(n),准确的讲是 O(4n)。当然,如果使用链式的方式存储线段树,则不需要 4n 的空间也可以。
应用
线段树常用于基于区间进行的统计查询。如果我们经常需要查询数组中某区间的最大值,最小值或者这个区间的数字和等,就可以使用线段树使操作复杂度由 O(n) 下降为 O(logn)。
具体的说,比如一个网站需要统计其2017年注册用户中至今为止消费最高的用户?消费最少的用户?学习时长最长的用户?
实现
public interface Merger<E> {
E merge(E a, E b);
}
public class SegmentTree<E> {
private E[] tree;
private E[] data;
// 融合器,区间中的元素进行何种融合存到线段树中。
private Merger<E> merger;
/**
* 构造函数。
*
* @param arr 传入的数组。
* @param merger 融合器。
*/
@SuppressWarnings("unchecked")
public SegmentTree(E[] arr, Merger<E> merger) {
this.merger = merger;
data = (E[]) new Object[arr.length];
for (int i = 0; i < arr.length; i++) {
data[i] = arr[i];
}
// 线段树所使用数组空间长度应该是传入数组长度的 4 倍
tree = (E[]) new Object[4 * arr.length];
buildSegmentTree(0, 0, data.length - 1);
}
/**
* 在 treeIndex 的位置创建表示区间[l...r]的线段树
*
* @param treeIndex 创建的线段树根节点所在的索引
* @param l 区间的左边
* @param r 区间的右边
*/
private void buildSegmentTree(int treeIndex, int l, int r) {
if (l == r) {
tree[treeIndex] = data[l];
return;
}
int leftTreeIndex = leftChild(treeIndex);
int rightTreeIndex = rightChild(treeIndex);
int mid = l + (r - l) / 2;
buildSegmentTree(leftTreeIndex, l, mid);
buildSegmentTree(rightTreeIndex, mid + 1, r);
tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
}
// 区间长度。
public int getSize() {
return data.length;
}
// 获得指定索引的元素。
public E get(int index) {
if (index < 0 || index >= data.length) {
throw new IllegalArgumentException("Index is illegal.");
}
return data[index];
}
// 左孩子索引。
private int leftChild(int index) {
return 2 * index + 1;
}
// 右孩子索引。
private int rightChild(int index) {
return 2 * index + 2;
}
// 返回区间 [queryL, queryR]的值
public E query(int queryL, int queryR) {
if (queryL < 0 || queryL >= data.length || queryR < 0 || queryR >= data.length || queryL > queryR) {
throw new IllegalArgumentException("Index is illegal.");
}
return query(0, 0, data.length - 1, queryL, queryR);
}
// 在以treeIndex 为根的线段树中[l...r]的范围内,搜索区间 [queryL...queryR]的值
private E query(int treeIndex, int l, int r, int queryL, int queryR) {
if (l == queryL && r == queryR) {
return tree[treeIndex];
}
int mid = l + (r - l) / 2;
int leftTreeIndex = leftChild(treeIndex);
int rightTreeIndex = rightChild(treeIndex);
if (queryL >= mid + 1) {
return query(rightTreeIndex, mid + 1, r, queryL, queryR);
} else if (queryR <= mid) {
return query(leftTreeIndex, l, mid, queryL, queryR);
}
E leftResult = query(leftTreeIndex, l, mid, queryL, mid);
E rightResult = query(rightTreeIndex, mid + 1, r, mid + 1, queryR);
return merger.merge(leftResult, rightResult);
}
// 将index位置的值,更新为e
public void set(int index, E e) {
if (index < 0 || index >= data.length) {
throw new IllegalArgumentException("Index is illegal");
}
data[index] = e;
set(0, 0, data.length - 1, index, e);
}
// 在以treeIndex为根的线段树中更新index的值为e
private void set(int treeIndex, int l, int r, int index, E e) {
if (l == r) {
tree[treeIndex] = e;
return;
}
int mid = l + (r - l) / 2;
// treeIndex的节点分为[l...mid]和[mid+1...r]两部分
int leftTreeIndex = leftChild(treeIndex);
int rightTreeIndex = rightChild(treeIndex);
if (index >= mid + 1) {
set(rightTreeIndex, mid + 1, r, index, e);
} else { // index <= mid
set(leftTreeIndex, l, mid, index, e);
}
tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
}
@Override
public String toString() {
StringBuilder res = new StringBuilder();
res.append('[');
for (int i = 0; i < tree.length; i++) {
if (tree[i] != null) {
res.append(tree[i]);
} else {
res.append("null");
}
if (i != tree.length - 1) {
res.append(", ");
}
}
res.append(']');
return res.toString();
}
}
使用线段树,可以以 O(logn) 的复杂度方便的进行数组区间求和、求区间最大值等操作:
public class Main {
public static void main(String[] args) {
Integer[] nums = { -2, 0, 3, -5, 2, -1 };
SegmentTree<Integer> segTree = new SegmentTree<>(nums, new Merger<Integer>() {
@Override
// 用于求和的线段树。
public Integer merge(Integer a, Integer b) {
return a + b;
// return Math.max(a, b); // 用于求最大值的线段树。
}
});
// SegmentTree<Integer> segTree = new SegmentTree<>(nums, (a, b) -> a + b); // Lambda 表达式 的写法
System.out.println(segTree.query(0, 2)); // 1
System.out.println(segTree.query(2, 5)); // -1
System.out.println(segTree.query(0, 5)); // -3
System.out.println(segTree);
}
}
更多
- 我们这里实现的线段树只能实现单个元素更新,没有实现区间更新。比如说,我们希望区间 [2,5]中所有元素 +3,在线段树中,需要将这个叶子节点以及他们的父节点都进行更新,复杂度会变为 O(n) 级别,比较慢。一个方式是进行懒惰更新,lazy 更新,我们只把线段树中具体到相应区间的节点进行更新,其下面的子节点先不进行更新,而使用一个 lazy 数组记录未更新的内容,之后如果进行查询时先查找 lazy 数组看是否有未更新的内容,如果有再更新一下。
- 线段树不仅仅适用于一维,二维甚至三维的区间问题都可以使用线段树解决。
- 关于区间操作,还有另外一个重要的数据结构:树状数组。