「线段树」第 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);
}