我之前的文章讲过fenwick树,这是一种高效计算前缀和的树。但是现实中不仅仅有计算前缀和的需求,还有范围查询需求。比如说求第3项到第6项的和,当然用fenwick树也能实现,可以用第六项前缀和减去第三项前缀和。但是这肯定不是最好的方案啊。所以这里我要介绍的是线段树segment tree,有时候也简称为段树。当然段树不仅仅是一维的,也有解决二维问题、甚至更高维度的问题的段树,但是这里我只写一维段树。
段树有点像二叉堆,逻辑上是一棵排序二叉树,但是物理结构是一个数组。先说逻辑结构吧。逻辑结构是把数组按照二分法进行拆分。如下图所示:
图中是一个长度为6的数组的段树。这个段树不是求和,而是求最大值。可以看到,原始数组数据都保存在叶子上。叶子表示长度为1的范围。而中间节点是叶子的范围合并。段树只有三个方法:构建、查询与修改。
段树的构建方法,可以使用递归构建。从根节点开始,根结点存储0~n-1的统计数据。然后按照二叉堆的方式递归构建。使用二分法,小的一半是左孩子,大的一半是右孩子。
段树的统计方法,需要区分三种场景:第一种场景是直接位于某个节点,那就直接返回查询就可以了。第二种场景是在节点的左孩子或右孩子上,这个时候递归进去查询就可以了。最复杂的是第三种场景,一半在左孩子,一半在右孩子上,这时需要查询两边,进行一次计算。
段树的修改方法与构建方法类似,也是递归修改。整体代码比较简单,我贴一下代码:
package com.youngthing.trees.segment;
import java.util.Arrays;
import java.util.Deque;
import java.util.LinkedList;
import java.util.function.BiFunction;
import java.util.function.Consumer;
/**
* 线段树
*/
public class SegmentTree<T> {
private final Object[] array;
private final int n;
private final BiFunction<T, T, T> function;
public SegmentTree(T[] input, BiFunction<T, T, T> function) {
this.n = input.length;
this.function = function;
int x = n;
while (Integer.bitCount(x) > 1) {
x += (x & -x);
}
array = new Object[x << 1];
// 类似堆的初始化方法
build(input, 0, 0, n - 1);
}
public Object[] getArray() {
return array;
}
public T get(int l, int r) {
// 使用递归
return get(0, 0, this.n - 1, l, r);
}
public void update(int index, T value) {
update(0, 0, n - 1, index, value);
}
private void build(T[] input, int v, int tl, int tr) {
if (tl == tr) {
array[v] = input[tl];
} else {
int tm = (tl + tr) / 2;
build(input, 2 * v + 1, tl, tm);
build(input, 2 * v + 2, tm + 1, tr);
array[v] = function.apply((T) array[2 * v + 1], (T) array[2 * v + 2]);
}
}
private T get(int i, int tl, int tr, int l, int r) {
if (tl == l && tr == r) {
return (T) this.array[i];
}
int tm = (tl + tr) / 2;
// 现在分两种场景:1 两边之和
if (tm >= l && tm < r) {
return function.apply(get(2 * i + 1, tl, tm, l, tm),
get(2 * i + 2, tm + 1, tr, tm + 1, r));
}
// 第二种场景 只在一边
// 只在左边
if (r <= tm) {
return get(2 * i + 1, tl, tm, l, r);
}
// 只在右边
return get(2 * i + 2, tm + 1, tr, l, r);
}
private void update(int i, int tl, int tr, int index, T value) {
if (tl == tr) {
array[i] = value;
} else {
int tm = (tl + tr) / 2;
if (index <= tm) {
update(2 * i + 1, tl, tm, index, value);
} else {
update(2 * i + 2, tm + 1, tr, index, value);
}
array[i] = function.apply((T) array[2 * i + 1], (T) array[2 * i + 2]);
}
}
public String toGraphString() {
StringBuilder builder = new StringBuilder();
builder.append("digraph segment{\nlayout=fdp\n");
// node
// bfs
int[] xCoordinates = new int[array.length];
int[] yCoordinates = new int[array.length];
int[] xLeaf = {0};
postOrder(x -> {
int i = x[0];
Object v = array[i];
// 有孩子
if (2 * i + 2 < array.length && array[2 * i + 2] != null) {
xCoordinates[i] = (xCoordinates[2 * i + 2] + xCoordinates[2 * i + 1]) / 2;
yCoordinates[i] = Integer.max(yCoordinates[2 * i + 1], yCoordinates[2 * i + 2]) + 1;
} else {
xCoordinates[i] = xLeaf[0] += 2;
yCoordinates[i] = 0;
}
builder.append("t").append(i).append("[shape=record;label=\"{{v=")
.append(v).append("}|{[")
.append(x[1]).append("~").append(x[2])
.append("]}}\";pos=\"").append(xCoordinates[i]).append(",")
.append(yCoordinates[i]).append("!\"]\n");
});
dfs(x -> {
int i = x[0];
int p = (i - 1) / 2;
if (p >= 0 && i != 0) {
builder.append("t").append(p).append("->")
.append("t").append(i).append("\n");
}
});
builder.append("}\n");
return builder.toString();
}
private void dfs(Consumer<Integer[]> consumer) {
Deque<Integer[]> queue = new LinkedList<>();
queue.add(new Integer[]{0, 0, n - 1});
while (!queue.isEmpty()) {
Integer[] pop = queue.pop();
int i = pop[0];
Object v = array[i];
if (v == null) {
continue;
}
consumer.accept(pop);
int tm = (pop[1] + pop[2]) / 2;
if (2 * i + 1 < array.length) {
queue.push(new Integer[]{2 * i + 1, pop[1], tm});
}
if (2 * i + 2 < array.length) {
queue.push(new Integer[]{2 * i + 2, tm + 1, pop[2]});
}
}
}
private void postOrder(Consumer<Integer[]> consumer) {
Deque<Integer[]> queue = new LinkedList<>();
dfs(queue::push);
while (!queue.isEmpty()) {
consumer.accept(queue.pop());
}
}
}
测试类:
public class SegmentTreeTest {
@Test
public void test() {
int[] ints = {1, 2, 3, 4, 5, 6};
Integer[] array = Arrays.stream(ints).mapToObj(Integer::valueOf).toArray(Integer[]::new);
SegmentTree<Integer> tree = new SegmentTree<>(array, Math::max);
System.out.println(Arrays.toString(tree.getArray()));
System.out.println(tree.get(2,4));
tree.update(3,7);
System.out.println(tree.get(0,5));
}
}
测试结果:
[6, 3, 6, 2, 3, 5, 6, 1, 2, null, null, 4, 5, null, null, null]
5
7