数据结构之线段树及其应用

线段树Segment Tree

线段树作为一种高级数据结构主要解决的是和区间有关的问题，当我们关注的问题为区间的某个统计量时(如和、最值等)，往往使用线段树这种数据结构。当然使用线性结构也可以达到这个目的，不过一般的线性结构来完成更新区间和查询区间操作的时间复杂度都是O(n)。

完全二叉树存储线段树示意图

而线段树是基于平衡二叉树的结构，每一次修改、查询的时间复杂度都为O(logn)，这是远远优于线性结构的。

线段树详解

对于线段树树结构的实现，为了简单起见，我们使用数组，有读者可能会问，数组不是要用来表示完全二叉树才有规律吗？没错，所以我们也将线段树存储为完全二叉树，虽然这样浪费了一些空间，但是能带来很大的方便。于是对于一个有n个元素的整体，要将它存储为线段树，4n个空间可保证足够。而且由于线段树一般不会考虑添加元素，区间固定，所以使用静态空间即可。

线段树的查询和更新

对于线段树我们只考虑查询和更新两个操作，由于树结构有天然的递归结构，所以这两个操作在逻辑上并不复杂，于是给出线段树类：

'''Merger.java'''
public interface Merger<E> {
    E merge(E a, E b);
}

'''SegmentTree.java'''
public class SegmentTree<E> {

    private E[] tree;
    private E[] data;
    private Merger<E> merger;

    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];

        tree = (E[])new Object[arr.length * 4];//4倍空间
        buildSegmentTree(0,0,data.length-1);
        //递归创建线段树
    }
    private void buildSegmentTree(int treeIndex, int l, int r){
        //在treeIndex位置创建表示区间[l,...,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 error!");
        return data[index];
    }

    //辅助函数，找左右孩子
    private int leftChild(int index){
        return 2 * index +1;
    }
    private int rightChild(int index){
        return 2 * index +2;
    }


    public E query(int queryL, int queryR){
        if(queryL < 0 || queryR >= data.length || queryL > queryR || queryR < 0 || queryL >= data.length)
            throw new IllegalArgumentException("Index is illegal!");
        return query(0,0,data.length-1, queryL, queryR);
    }
    private E query(int treeIndex, int l, int r, int queryL, int queryR){
        //在以treeIndex为根的线段树[l,..,r]，搜索区间[queryL,..,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);
    }
    public void set(int index, E e){
        //线段树更新操作
        if(index <0 || index >= data.length)
            throw new IllegalArgumentException("Index is error!");
        data[index] = e;
        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 mid = l + (r-l)/2;
        int leftTreeIndex = leftChild(treeIndex);
        int rightTreeIndex = rightChild(treeIndex);
        if(index >= mid + 1)
            set(rightTreeIndex, mid + 1, r, index, e);
        else
            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();
    }
}

说明：

• Merger类是用于用户定义合并子树规则的。
• 构造函数接收用户输入的一个E（泛型）型数组，将其拷贝到data数组之后，通过data中的数据和buildSegmentTree函数以递归方式构建一棵线段树tree。
• 这里线段树的左半边子树长度是其父亲长度除以2后向下取整得到。
• 在线段树的更新操作（set）中，传入的前三个参数是分别是线段树的根、根对应线段树元素下标的始点和下标的终点，实际上这三个参数可以封装为一个，不过这里为了清晰展示逻辑没有这样做。
• 注意更新操作会造成祖先节点值的改变，所以更新操作不仅仅是修改数组的值，一个值被修改则祖先都会改变。
• 线段树合并操作（merge）是根据具体的业务逻辑决定的，可以传入我们的规则进去，比如我们可以规定为求和、求最大值等等。

---

线段树的应用——LeetCode303、307

先来看一道简单的题目：

LeetCode303

这里要求任意子列的和，回想我们所学的线段树，正好可以解决这类问题，于是给出解答：

class NumArray {
    //这里要添加前面写的接口和线段树作为内部类再提交，为节约篇幅，这里仅文字交代下
    private SegmentTree<Integer> segmentTree;
    public NumArray(int[] nums) {
        if(nums.length > 0){
            Integer[] data = new Integer[nums.length];
            for(int i=0; i < nums.length; i++)
                data[i] = nums[i];
            segmentTree = new SegmentTree<>(data, (a, b) -> a + b);
        }
    }

    public int sumRange(int i, int j) {
        if(segmentTree == null)
            throw new IllegalArgumentException("SegmentTree is null!");
        return segmentTree.query(i, j);
    }
}

其中的合并规则定义为相加，提交，获得通过！不过有的读者应该发现了，对于本题其实不用线段树直接用循环对数组子列求和也是可以的，甚至代码更简洁，具体如下：

class NumArray {//不使用线段树,更加高效

    private int[] sum;//sum[i]存储前i个元素的和,sum[0]=0，nums[0,...,i-1]的和
    public NumArray(int[] nums){
        sum = new int[nums.length + 1];
        sum[0] = 0;
        for(int i=1;i<sum.length;i++)
            sum[i] = sum[i-1] + nums[i-1];
    }
    public int sumRange(int i, int j){
        return sum[j+1] - sum[i];
    }
}

提交，同样会获得通过。既然不用线段树也可以，那为什么还要研究线段树呢？这就涉及到线段树的性能优势。下面来看这道题的进阶版：

LeetCode307

与上一题不同的是，同样是求子列和，这里的数组中的元素可以被修改，于是按照上题的逻辑给出不用线段树的解答：

class NumArray {//普通方案在update操作时最坏时间复杂度O(n)，逻辑正确但提交可能会超时

    private int[] sum;//sum[i]存储前i个元素的和,sum[0]=0，nums[0,...,i-1]的和
    private int[] data;
    public NumArray(int[] nums){
        data = new int[nums.length];
        for(int i=0;i<nums.length;i++)
            data[i] = nums[i];
        sum = new int[nums.length + 1];
        sum[0] = 0;
        for(int i=1;i<sum.length;i++)
            sum[i] = sum[i-1] + nums[i-1];
    }
    public void update(int index, int val) {//最坏O(n)
        data[index] = val;
        for(int i=index+1;i<sum.length;i++)
            sum[i] = sum[i-1] + data[i-1];
    }
    public int sumRange(int i, int j){
        return sum[j+1] - sum[i];
    }
}

提交后发现并没有通过，然而算法的逻辑是没有问题的。这是因为LeetCode对一些题目的解答有时间的限制，而修改操作是比较复杂的。对于线性结构来说，修改一个元素的值后，还要一步步修改其祖先节点，这代价是巨大的，于是换用线段树：

public class NumArray2 {//使用线段树
    //别忘了提交时带上内部类：线段树类和合并器接口
    private SegmentTree<Integer> segmentTree;
    public NumArray2(int[] nums) {
        if(nums.length > 0){
            Integer[] data = new Integer[nums.length];
            for(int i=0; i < nums.length; i++)
                data[i] = nums[i];
            segmentTree = new SegmentTree<>(data, (a, b) -> a + b);
        }
    }

    public int sumRange(int i, int j) {
        if(segmentTree == null)
            throw new IllegalArgumentException("SegmentTree is null!");
        return segmentTree.query(i, j);
    }
    public void update(int index, int val) {//最坏O(n)
        if(segmentTree == null)
            throw new IllegalArgumentException("segmentTree is null");
        segmentTree.set(index,val);
    }

}

再提交，获得通过！

通过

小结

线段树是一种很有用的数据结构，经典的区间染色问题可以很轻松用线段树解决，总之如果我们关注的是和子区间相关的问题，而且区间具有可加性时都可以考虑使用线段树这种数据结构。