6.12 线段树

  我之前的文章讲过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