「線段樹」第 3 節：創建線段樹與區間查詢

根據原始數組創建線段樹

這一節的目標是：我們把員工的信息輸入一棵線段樹，讓這棵線段樹組織出領導架構。即已知 data 數組，要把 tree 數組構建出來。

• 分析遞歸結構，重點體會：二叉樹每做一次分支都是「一分爲二」進行的，因此線段樹是一棵二叉樹；
• 遞歸到底的時候，這個區間只有 $1$ 個元素。

設計私有函數，我們需要考慮：

• 我們要創建的線段樹的根結點的下標，這個下標是線段樹的下標；
• 對於線段樹結點所要表示的 data 數組的區間的左端點是什麼；
• 對於線段樹結點所要表示的 data 數組的區間的右端點是什麼。

Java 代碼：

buildSegmentTree(0, 0, arr.length - 1);

Java 代碼：關鍵代碼

/**
 * 這個遞歸方法的描述一定要非常清楚：
 * 畫出 tree 樹中以 treeIndex 爲根的，統計 data 數組中 [l,r] 區間中的元素
 * 這個方法的實現引入了一個 merge 接口，使得外部可以傳入一個方法，方法是如何實現的是根據業務而定
 * 核心代碼只有幾行，這裏關鍵還是在於遞歸方法
 *
 * @param treeIndex 我們要創建的線段樹根結點所在的索引，treeIndex 是 tree 的索引
 * @param l         對於 treeIndex 結點所要表示的 data 區間端點是什麼，l 是 data 的索引
 * @param r         對於 treeIndex 結點所要表示的 data 區間端點是什麼，r 是 data 的索引
 */
private void buildSegmentTree(int treeIndex, int l, int r) {
    // 考慮遞歸到底的情況
    if (l == r) {
        // 平衡二叉樹葉子結點的賦值就是靠這句話形成的
        tree[treeIndex] = data[l]; // data[r]，此時對應葉子結點的情況
        return;// return 不能忘記
    }
    int mid = l + (r - l) / 2;
    int leftChild = leftChild(treeIndex);
    int rightChild = rightChild(treeIndex);
    // 假設左邊右邊都處理完了以後，再處理自己
    // 這一點基於，高層信息的構建依賴底層信息的構建
    // 這個遞歸的過程我們可以通過畫圖來理解
    // 仔細閱讀下面的這三行代碼，是不是像極了二分搜索樹的後序遍歷，我們先處理了左右孩子結點，最後處理自己
    buildSegmentTree(leftChild, l, mid);
    buildSegmentTree(rightChild, mid + 1, r);
    
    // 注意：merge 的實現根據業務而定
    tree[treeIndex] = merge.merge(tree[leftChild], tree[rightChild]);
}

Merge 接口的設計，這裏使用傳入對象的方式實現了方法傳遞，是 Command 設計模式。
Java 代碼：

public interface Merge<E> {
    E merge(E e1, E e2);
}

給 SegmentTree 覆蓋 toString 方法，用於打印線段樹表示的數組，以便執行測試用例。

Java 代碼：

@Override
public String toString() {
    StringBuilder s = new StringBuilder();
    s.append("[");
    for (int i = 0; i < tree.length; i++) {
        if(tree[i] == null){
            s.append("NULL");
        }else{
            s.append(tree[i]);
        }
        s.append(",");
    }
    s.append("]");
    return s.toString();
}

測試方法：

public class Main {
    public static void main(String[] args) {
        Integer[] nums = {0, -1, 2, 4, 2};
        SegmentTree<Integer> segmentTree = new SegmentTree<Integer>(nums, new Merge<Integer>() {
            @Override
            public Integer merge(Integer e1, Integer e2) {
                return e1 + e2;
            }
        });
        System.out.println(segmentTree);
    }
}

區間查詢

通過編寫二分搜索樹的經驗，我們知道，一些遞歸的寫法通常要寫一個輔助函數，在這個輔助函數裏完成遞歸調用。那麼對於這個問題中，輔助函數的設計就顯得很關鍵了。

// 在一棵子樹裏做區間查詢，dataL 和 dataR 都是原始數組的索引
public E query(int dataL, int dataR) {
    if (dataL < 0 || dataL >= data.length || dataR < 0 || dataR >= data.length || dataL > dataR) {
        throw new IllegalArgumentException("Index is illegal.");
    }
    // data.length - 1 邊界不能弄錯
    return query(0, 0, data.length - 1, dataL, dataR);
}

在這個輔助函數的實現過程中，可以畫一張圖來展現一下具體的計算過程。

體會下面這個過程：我們總是自上而下，從根結點開始向下查詢，最壞情況下，纔會查詢到葉子結點。
Java 代碼：

// 這是一個遞歸調用的輔助方法，應該定義成私有方法
private E query(int treeIndex, int l, int r, int dataL, int dataR) {
    if (l == dataL && r == dataR) {
        // 這裏一定不要犯暈，看圖說話
        return tree[treeIndex];
    }
    int mid = l + (r - l) / 2;
    int leftChildIndex = leftChild(treeIndex);
    int rightChildIndex = rightChild(treeIndex);
    // 畫個示意圖就能清楚自己的邏輯是怎樣的
    if (dataR <= mid) {
        return query(leftChildIndex, l, mid, dataL, dataR);
    }
    if (dataL >= mid + 1) {
        return query(rightChildIndex, mid + 1, r, dataL, dataR);
    }
    // 橫跨兩邊的時候，先算算左邊，再算算右邊
    E leftResult = query(leftChildIndex, l, mid, dataL, mid);
    E rightResult = query(rightChildIndex, mid + 1, r, mid + 1, dataR);
    return merge.merge(leftResult, rightResult);
}