数据结构与算法-线段树
图片来自慕课网，liuyubobobo讲师的课程“玩转数据结构 从入门到进阶”
线段树介绍
有时候需要对某个区间进行操作，比如求和、求最大值最/小值等。下面的问题：
有数组arr = [2,6,3,5,7,1]，求数组任意[L, R]区间上的和。
最直观的方法就是对[L, R]区间遍历求和，最差情况时间复杂度O(N)。使用线段树可以达到O(lg n)的时间复杂度。
什么是线段树？线段树又叫区间树，将某个区间用一个树结点表示，树可以像二叉堆那样用数组在构建。如下图所示数组A的长度是8，根结点表示整个区间，即[0, 7]；根结点的左孩子表示区间的左半部分[0, 3]，右孩子表示区间的右半部分[4, 7]。根据二叉树的递归结构，对于A[0, 3]和A[4, 7]像上面一样继续构建。当达到树的叶子结点时，区间中只有一个元素。
更广泛的说，线段树的每个结点都表示数组A的[L, R]区间，在树的叶子结点处，L == R，是单元素区间，表示数组中的某一个元素。
我们要找某个[L, R]区间上的某种值（最大、最小、和），只需要从对应的树结点上获取即可。比如求A[2, 3]的和，只需要像二叉搜索树一样，从根结点出发、选择左孩子A[0, 3]再选择右孩子A[2, 3]返回其值即可。但是树结点表示的区间并不能囊括所有的区间可能。比如现在要求A[2, 5]的和。其实可以看作是A[2, 3]和A[4, 5]两个结点的结点的和，也就是结果分散在树中的多个结点中；或者说可以将区间拆分成多个子区间，不断拆分，直到能在树中找到对应的区间，然后将找到的所有结点的结果汇总起来即可。因此，要求A[2, 6]其实就是将A[2, 3]、A[4, 5]、A[6]的和汇总起来。
上面元素个数是8，树的最后一层刚好可以容纳下所有的元素，此时树是一棵满二叉树。但是如果元素个数有10个的话，这一层就容纳不下了，因此需要再增加一层，如下，此时树不是满二叉树，甚至完全二叉树也不是。但是线段树是平衡二叉树。
为了将树表示成满二叉树，可以将null结点也补上。
我们还看到对于奇数个元素的数组，区间不能平均分，可以固定让左孩子的区间比右孩子的区间小。如上图中的A[2]和A[3, 4]。
那么对于n个元素的数组，需要多大的空间（包含补上的null结点）？
根据满二叉树的性质，树的最后一层结点个数是，而树总的结点树是，可以认为树的最后一层结点个数是树总结点的一半。
因为我们不考虑往线段树中添加元素，因此使用4n大小的静态数组即可，虽然大多数情况下会有很多null结点浪费了数组空间，但是能保证在最坏情况下也不至于数组下标越界。
线段树的实现
根据上面对线段树的描述，树的叶子结点存放的是数组中的每一个元素。我们需要一个data数组作为传入数组的一个副本，同时用一个tree数组表示线段树。树的根结点使用tree[0],表示了数组的整个区间。对于任何一个非叶子结点i,其左孩子在数组中的2 * i + 1处，右孩子在数组中的2 * i + 2处。
线段树的实现骨架如下
package tree;
import java.util.Arrays;
public class SegmentTree {
/**
* 存放数据的数组
*/
private E[] data;
/**
* 线段树，使用4n的静态空间
*/
private E[] tree;
private int size;
public int size() {
return size;
}
public E get(int index) {
if (index < 0 || index >= size) {
throw new ArrayIndexOutOfBoundsException(index);
}
return data[index];
}
@Override
public String toString() {
StringBuilder sb = new StringBuilder();
sb.append("[");
for (int i = 0; i < tree.length; i++) {
sb.append(tree[i]);
if (i != tree.length - 1) sb.append(", ");
}
return sb.append("]").toString();
}
}
get方法简单的返回传入数组中index处的元素而已。
现在要将传入的数组转换成线段树，以方便对任何区间进行操作。这个操作可以是求和，求最大/最小等等。不如实现一个接口，让用户自己实现要完成的操作。如下Merger接口，对两个元素操作返回一个结果。
package tree;
public interface Merger {
E merge(E a, E b);
}
来看构造方法，data是传入的数组，Merger就是我们要进行的操作。
public SegmentTree(E[] data, Merger merger) {
this.size = data.length;
this.merger = merger;
this.data = Arrays.copyOf(data, data.length);
this.tree = (E[]) new Object[4 * data.length];
buildSegmentTree(0, 0, data.length - 1);
}
比如要求区间和，匿名内部类就可以写成下面这样
new Merger() {
@Override
public Integer merge(Integer a, Integer b) {
return a + b;
}
});
线段树的构建
构造方法中的buildSegmentTree(treeIndex, L, R);方法是关键。从根结点开始，递归地构建线段树：
- 递归终止条件是到树的叶子结点。即区间的L和R相等时，此时区间中只有一个元素，将
data[l]或者data[r]值赋给treeIndex处的结点即可，表示treeIndex处的结点保存着该单元素区间的值。 - 否则，递归的构建左右子树：先找到当前结点的左、右孩子treeIndex，然后将区间
[l, r]平分成两部分，左子树leftChild表示区间[l, mid]，右子树rightChild表示区间[mid + 1, r]，在这两个区间递归的构建左右子树 - 左右子树构建好了之后，可以利用左右孩子的区间值汇总成根结点的值（树的后序遍历思想），因为根结点的区间正是左右孩子表示区间的结合。
代码如下
/**
* 构建线段树
* @param treeIndex SegmentTree中以treeIndex为根的树
* @param l treeIndex对应的树的左区间
* @param r treeIndex对应的树的右区间
*/
private void buildSegmentTree(int treeIndex, int l, int r) {
// 叶子结点，左右区间一样，index处存放的数据是data[l]或者data[r]
if (l == r) {
tree[treeIndex] = data[l];
return;
}
int leftChild = 2 * treeIndex + 1;
int rightChild = leftChild + 1;
int mid = (r - l) / 2 + l;
buildSegmentTree(leftChild, l, mid);
buildSegmentTree(rightChild, mid + 1, r);
// 左右树建立好后可以针对具体业务场景，将左右树结果“汇总”到当前树结点中
// 如对区间求和那么 tree[treeIndex] = tree[leftChild] + tree[rightChild]
// 再如求区间的最大/最小值 tree[treeIndex] = Math.max(tree[leftChild], tree[rightChild])
tree[treeIndex] = merger.merge(tree[leftChild], tree[rightChild]);
}
线段树的区间查询
这是核心操作，用于查询区间[L, R]上的某操作后的值。
public E query(int queryL, int queryR) {
if (queryL < 0 || queryR < 0 || queryL >= data.length || queryR >= data.length || queryL > queryR) {
throw new IllegalArgumentException("index is illegal");
}
return query(0, 0, data.length - 1, queryL, queryR);
}
/**
* 查询[queryL, queryR]之间的结果
* @param treeIndex SegmentTree中以treeIndex为根的树
* @param l treeIndex对应的树的左区间
* @param r treeIndex对应的树的右区间
* @param queryL 查找的左范围
* @param queryR 查找的右范围
* @return [queryL, queryR]之间的结果
*/
private E query(int treeIndex, int l, int r, int queryL, int queryR) {
// 如果查找的区间刚好树的区间对应上了，直接返回treeIndex处的结果
if (l == queryL && r == queryR) {
return tree[treeIndex];
}
int leftChild = 2 * treeIndex + 1;
int rightChild = leftChild + 1;
int mid = (r - l) / 2 + l;
// 如果查找的左范围比mid还大，只需要在右子树中查找
if (queryL > mid) return query(rightChild, mid + 1, r, queryL, queryR);
// 如果查找的右范围比mid + 1小，只需要在左子树中查找
if (queryR < mid + 1) return query(leftChild, l, mid, queryL, queryR);
// 在左右子树查找并将结果融合，将查询的区间[queryL，queryR]拆分成：在左子树中查找[queryL, mid]，在右子树中查找[mid + 1, queryR]
return merger.merge(query(leftChild, l, mid, queryL, mid), query(rightChild, mid + 1, r, mid + 1, queryR));
}
如果要查询的区间[queryL, queryR]刚好和树中某个结点表示的区间一样，便可以直接返回。否则在其左、右子树中查询，如果要查询的区间左端点比左子树的右端点还大，那么无需查询左子树，直接返回右子树的查询结果即可；类似的，如果要查询的区间右端点比右子树的左端点还小，就不必在右子树中查询了，直接返回左子树的查询结果即可；否则，要查询的区间覆盖了左右子树，需要将[queryL, queryR]拆分成[queryL,mid]和[]需要将左右子树的结果汇总。
线段树的更新
线段树的更新将除了index处元素更新成e，由于data数组改变，tree也需要更新。和构建线段树一样， L == R时在tree中找到了存放data[index]的结点，更新并返回；否则，如果查找的index比mid大，说明要更新的结点在右子树中，如果比mid + 1小说明要更新的结点在左子树中。和在buildSegmentTree中一样，在set方法中更新了叶子结点的值后，其上直到根结点的父结点都需要更新。
public void set(int index, E e) {
if (index < 0 || index >= size) {
throw new ArrayIndexOutOfBoundsException(index);
}
data[index] = e;
// 接下来更新tree,从根结点开始
set(0, 0, data.length - 1,index, e);
}
private void set(int treeIndex,int l, int r, int index, E e) {
if (l == r) {
tree[treeIndex] = e;
return;
}
int leftChild = 2 * treeIndex + 1;
int rightChild = leftChild + 1;
int mid = (r - l) / 2 + l;
if (index > mid) set(rightChild, mid + 1, r, index, e);
if (index < mid + 1) set(leftChild, l, mid, index, e);
tree[treeIndex] = merger.merge(tree[leftChild], tree[rightChild]);
}
测试
如下，对于数组[2, 3, 5, 6]，求[1, 3]区间上的和。
public static void main(String[] args) {
Integer[] data = {2, 3, 5, 6};
SegmentTree segmentTree = new SegmentTree<>(data, (a, b) -> a + b);
System.out.println(segmentTree.query(1,3)); // 14
segmentTree.set(2, 10);
System.out.println(segmentTree.query(1,3)); // 19
}
set前和是14，set后数组变成[2, 3, 10, 6]再查询结果是19。
@sunhaiyu
2018.11.24