【数据结构】数组区间更新-线段树

在讲解算法原型之前，我们先来看一道算法题，这道算法题很贴近生活，也就是我们小时候玩的俄罗斯方块。LeetCode699掉落的方块

题目描述的文字太多，我就简单点说，类似于俄罗斯方块，从上方掉下方块，问你在每一个位置的最高的高度是多少？如图：

就类似于上图。

读了原题目的话，你可以知道，这所谓的“类似于俄罗斯方块”，实际上就是映射在数组上进行一些操作。

本期文章的参考代码以及LeetCode699：github

一、线段树的认识以及构造方法

线段树：就是为了在数组上的某一段范围内，进行快速的增删改查操作。

在进行数组下标的计算时，为了让这个线段树更容易实现，我们将数组的第一个元素舍弃掉，也就是说将原数组的数据全部向后面移动一个下标的位置，把第一个位置空出来。

说到这里后，我们就可以认识认识线段树，线段树就是一颗二叉树，具体线段树有哪些东西呢，如下。

• length：表示舍弃后的数组的长度，也就是原数组长度+1
• arr[]：存放原数组的数据向后移动一个下标的所有数据。（对应上图，舍弃后的数组）
• int[] sum：这个数组比较特殊，一般线段树是计算某一个范围内的总和，所以这里是sum，而LeetCode699题，是计算某一个范围内的最大值，也可以写成max。
• int[] lazy：线段树的核心部分，意思是：假设需要在arr的L~R范围内，每个位置增加10,；而lazy数组中，在rt位置就指向L和R范围，我们只需将10放入lazy数组的rt位置中，等后面需要向下操作时，才往下进行添加。（简单说，就是暂存，先暂时保存在lazy数组，什么时候需要了，就再往下面发出去）
• int[] change：线段树的核心部分，跟lazy数组比较相似，lazy数组对应的是add操作，而change数组则表示update操作。意思是将某一个范围内的数据，改为一个新的数据。（也是暂存）
• boolean[] update：布尔类型。是true，表示当前位置有更新操作；是false，表示当前位置没有更新操作

可能你会问为什么没有删除操作，其实也就可以用add操作进行代替了，取相反数即可。

对应的代码如下：

public class SegmentTree {
    private int length; //数组的长度再+1，因为0下标位置不需要
    private int[] arr; //原数组的所有数据，只是整体向右移动了一位
    private int[] sum; //某个范围上的总和
    private int[] lazy; //懒数组，为了实现某个范围上的操作，只需改这一个位置的数据
    private int[] change; //update更新的数据
    private boolean[] update; //表示相应的下标位置是否需要进行update
    
}

现在的问题就是，上面的每个数组，需要开辟多大的内存空间呢？

除了arr数组，是原数组的长度+1。其他的4个数组，都需要开辟 4 * arr.length。至于为什么是4倍，我们这里就不细说了。可以看看这个证明：https://blog.csdn.net/smoggyxhdz/article/details/78895672

所以代码如下：

public class SegmentTree {
    private int length; //数组的长度再+1，因为0下标位置不需要
    private int[] arr; //原数组的所有数据，只是整体向右移动了一位
    private int[] sum; //某个范围上的总和
    private int[] lazy; //懒数组，为了实现某个范围上的操作，只需改这一个位置的数据
    private int[] change; //update更新的数据
    private boolean[] update; //表示相应的下标位置是否需要进行update

    public SegmentTree(int[] arr) {
        length = arr.length + 1;
        this.arr = new int[length];
        for (int i = 1; i < length; i++) {
            this.arr[i] = arr[i - 1];
        }

        sum = new int[length << 2]; //4倍的长度
        lazy = new int[length << 2];
        change = new int[length << 2];
        update = new boolean[length << 2];
    }
}

二、build方法

build方法就是对sum数组进行初始化。这个方法有三个参数：左边界、右边界和sum数组的下标rt，这里说的左右边界指的是对应到arr数组的。比较简单，就是一个递归函数。

下标rt的计算，是根据二叉树的性质得到的。假设父节点的下标是i，则左孩子的下标就是 2 * i，右孩子的下标就是 2 * i+ 1。（舍弃0下标的情况下）

/**
 * 对sum数组进行初始化
 * @param l  左边界，最小是1
 * @param r  右边界，最大是arr数组的元素个数
 * @param rt sum的下标
 */
public void build(int l, int r, int rt) {
    if (l == r) {
        sum[rt] = arr[l];
        return;
    }
    //二分，进行递归
    int mid = (l + r) >> 1; //中位数
    build(l, mid, rt << 1); //左子树
    build(mid + 1, r, rt << 1 | 1); //右子树
    pushUp(rt); //两边汇总
}

这里写到了一个pushUp方法。这个方法就是将左右孩子的总和加起来，放入sum数组的rt位置。如下

private void pushUp(int rt) {
    sum[rt] = sum[rt << 1] + sum[rt << 1 | 1]; //当前节点的总和，就等于左右两个孩子的总和
}

这个方法是可以自己根据题意改的，比如上文提到的LeetCode699的题，是计算范围内的最大值，这里就可以这样写：

//LeetCode699 掉落的方块
private void pushUp(int rt) {
    max[rt] = Math.max(max[rt << 1], max[rt << 1 | 1]); //两边取最大值
}

顺带说一句，rt<<1也就是rt*2。rt<<1 | 1也就是rt * 2 + 1

三、add方法

我们最开始就说了线段树是为了解决数组中某一个范围内的操作。add方法就对某一个范围内进行操作，自然参数也就少不了左右边界。

• int L，表示指定范围的左边界。固定值
• int R，表示指定范围内的右边界。固定值
• int C，表示添加的数。固定值
• int curL，表示当前递归函数的左边界。可变参数
• int curR，表示当前递归函数的右边界。可变参数
• int rt，指向lazy数组的下标。

public void add(int L, int R, int C, int curL, int curR, int rt) {
    if (L <= curL && R >= curR) { //被加工的范围超过了当前递归的范围，可以懒
        sum[rt] += C * (curR - curL + 1); //当前位置的总和+ C*个数
        lazy[rt] += C; //懒数组的参数加C
        return;
    }
    //不能懒，直接再取中位数，递归调用
    int mid = (curL + rcurR) >> 1;
    //先将上次lazy数组剩下的参数往下发
    //mid -curL+1，左子树的节点数
    //curR-mid，右子树的节点数
    pushDown(rt, mid - curL + 1, curR- mid);
    if (L <= mid) {
        add(L, R, C, curL, mid, rt << 1); //递归左子树
    }
    if (R > mid) {
        add(L, R, C, mid + 1, curR, rt << 1 | 1); //递归右子树
    }
    pushUp(rt); //当前位置做汇总
}

递归的结束条件是当前递归的范围在指定范围LR内，就可以进行“偷懒”，将添加的参数放入lazy数组。

如上图，就是可以进行懒的情况。

而怎么进行懒？也很简单

//计算curL ~ curR范围内的总和
sum[rt] += C * (curR - curL + 1); //当前位置的总和+ C*个数
//将每个位置新增加的数放入lazy数组
lazy[rt] += C; //懒数组的参数加C

然后就是pushDown方法，与pushUp是相反的。pushUp是将左右孩子的总和加起来，而pushDown是将当前节点的数据分发给左右孩子。

pushDown方法是用于add方法和update方法，所有这个方法里面分为两个部分。一部分是针对add方法的，另外一部分是针对update方法的。这里我只先说针对add方法的部分。

/**
 * @param rt sum数组的下标
 * @param ln 左子树的节点数
 * @param rn 右子树的节点数
 */
private void pushDown(int rt, int ln, int rn) {
    //针对add方法的部分
    //懒数组，往下更新。也就是将当前数组的数据，往下分发
    if (lazy[rt] != 0) {
        lazy[rt << 1] += lazy[rt]; //左孩子
        lazy[rt << 1 | 1] += lazy[rt]; //右孩子
        
        sum[rt << 1] += ln * lazy[rt]; //左孩子的总和增加
        sum[rt << 1 | 1] += rn * lazy[rt]; //右孩子的总和增加
        lazy[rt] = 0; //lazy数组归0
    }
}

三、update方法

update，是对某一个范围内的所有位置，改成一个新的数据。递归的模板还是跟add方法一样，我们只需改动部分代码即可。

public void add(int L, int R, int C, int curL, int curR, int rt) {
    if (L <= curL && R >= curR) { //被加工的范围超过了当前递归的范围，可以懒
        update[rt] = true;
        change[rt] = C; //每个位置改的数
        sum[rt] = C * (curR - curL + 1); //新的总和
        lazy[rt] = 0; //lazy数组归0
        return;
    }
    //不能懒，直接再取中位数，递归调用
    int mid = (curL + rcurR) >> 1;
    //先将上次lazy数组剩下的参数往下发
    //mid -curL+1，左子树的节点数
    //curR-mid，右子树的节点数
    pushDown(rt, mid - curL + 1, curR- mid);
    if (L <= mid) {
        add(L, R, C, curL, mid, rt << 1); //递归左子树
    }
    if (R > mid) {
        add(L, R, C, mid + 1, curR, rt << 1 | 1); //递归右子树
    }
    pushUp(rt); //当前位置做汇总
}

递归的结束条件还是一样的。然后将update【rt】改为true，change【rt】位置改为新的数据，重新计算sum【rt】的总和即可，最后要记得将lazy【rt】位置改为0，也就是说以前add的所有都无效了。

现在再来说pushDown方法的另外一个针对update操作的部分代码。我们最开始就说了pushDown方法是将当前位置的数据，往左右孩子进行分发。这里也是一样的。

/**
 * @param rt sum数组的下标
 * @param ln 左子树的节点数
 * @param rn 右子树的节点数
 */
private void pushDown(int rt, int ln, int rn) {
    //针对update方法
    if (update[rt]) {
        update[rt << 1] = true; //将孩子节点改为true
        update[rt << 1 | 1] = true;
        
        change[rt << 1] = change[rt]; //孩子的change继承父节点
        change[rt << 1 | 1] = change[rt];
        
        lazy[rt << 1] = 0; //孩子的懒数组归0
        lazy[rt << 1 | 1] = 0;
        
        sum[rt << 1] = ln * change[rt]; //左孩子的总和
        sum[rt << 1 | 1] = rn * change[rt]; //右孩子的总和
        update[rt] = false; //当前位置已经往下发了，重新改为false
    }

    //针对add方法
    //懒数组，往下更新。也就是将当前数组的数据，下分发
    if (lazy[rt] != 0) {
        lazy[rt << 1] += lazy[rt];
        lazy[rt << 1 | 1] += lazy[rt];
        sum[rt << 1] += ln * lazy[rt]; //左孩子的总和增加
        sum[rt << 1 | 1] += rn * lazy[rt]; //右孩子的总和增加
        lazy[rt] = 0;
    }
}

跟update方法的递归结束条件差不多的，将当前位置的change数据发给孩子，孩子的lazy数组等都需要进行改动。

四、query方法

查询方法就比较简单的，又不需要进行任何的改动数组的操作，直接拿取sum数组的参数即可。

/**
 * 查询累加和
 * @param L  目标的左边界
 * @param R  目标的右边界
 * @param l  当前递归的左边界
 * @param r  当前递归的右边界
 * @param rt sum数组的下标
 * @return 返回总和
 */
public long query(int L, int R, int curL, int curR, int rt) {
    if (L <= curL && R >= curR) {
        return sum[rt]; //直接拿取
    }
    long res = 0;
    int mid = (curL + curR) >> 1;
    pushDown(rt, mid - curL + 1, curR - mid);
    if (L <= mid) {
        res += query(L, R, curL, mid, rt << 1); //左子树
    }
    if (R > mid) {
        res += query(L, R, mid + 1, curR, rt << 1 | 1); //右子树
    }
    return res;
}

五、如何调用该线段树

上面我们说了如何写一个线段树，这里我们在说一说如何用一个线段树

public static void main(String[] args) {
    int[] origin = { 2, 1, 1, 2, 3, 4, 5 };
    SegmentTree seg = new SegmentTree(origin);
    int start = 1; // 整个区间的开始位置，规定从1开始，不从0开始 -> 固定
    int end = origin.length; // 整个区间的结束位置，规定能到N，不是N-1 -> 固定
    int root = 1; // 整棵树的头节点位置，规定是1，不是0 -> 固定
    int L = 2; // 操作区间的开始位置 -> 可变
    int R = 5; // 操作区间的结束位置 -> 可变
    int C = 4; // 要加的数字或者要更新的数字 -> 可变
    // 区间生成，必须在[S,N]整个范围上build
    seg.build(start, end, root);
    // 区间修改，可以改变L、R和C的值，其他值不可改变
    seg.add(L, R, C, start, end, root);
    // 区间更新，可以改变L、R和C的值，其他值不可改变
    seg.update(L, R, C, start, end, root);
    // 区间查询，可以改变L和R的值，其他值不可改变
    long sum = seg.query(L, R, start, end, root);
    System.out.println(sum);
}

牢记线段树是对一段范围内的数据进行操作，这样的话，就能很好的理解了。剩下的就是多写，写熟练即可。

好啦，本期更新就到此结束啦！！！