數據結構之線段樹及其應用

線段樹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

先來看一道簡單的題目：

這裏要求任意子列的和，回想我們所學的線段樹，正好可以解決這類問題，於是給出解答：

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

提交，同樣會獲得通過。既然不用線段樹也可以，那爲什麼還要研究線段樹呢？這就涉及到線段樹的性能優勢。下面來看這道題的進階版：

與上一題不同的是，同樣是求子列和，這裏的數組中的元素可以被修改，於是按照上題的邏輯給出不用線段樹的解答：

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);
    }

}

再提交，獲得通過！

小結

線段樹是一種很有用的數據結構，經典的區間染色問題可以很輕鬆用線段樹解決，總之如果我們關注的是和子區間相關的問題，而且區間具有可加性時都可以考慮使用線段樹這種數據結構。