左神算法笔记: 2. 线段树(区间修改数)

algorithem/左神算法/基础

更快速的、区间的更改、查询一些东西
eg. 实现一种数据结构使得以下三种计算的时间复杂度为O(logN )
void add(arr, L, R, V); ：将arr数组中的第L到R的数都增加V
void update(arr, L, R, V);：将arr数组中的第L到R的数都修改V
int getSum(arr, L, R);：将arr数组中的第L到R的数相加
需要多少空间：

• 最好情况：恰好是N = 2^n个数字，例如8，需要2N-1个空间 < 2N的空间
• 最差情况：恰好是N = 2^n + 1个数字，例如5，需要 2 * ( 2N - 2 ) -1 = 4N - 5 < 4N的空间
  

::为什么从下标1开始::：这样i节点对应的两个子节点的坐标分别为 2 * i 和 2 * i + 1，可以用位运算表示为 i << 1 和 i << 1 | 1，位运算快

懒更新：不彻底下发任务，存到lazy里

public class SegmentTree {
    private int MAXN;
    private int[] arr;
    private int[] sum;
    private int[] lazy;
    private int[] change;
    private boolean[] update;

    public SegmentTree(int[] origin) {
        MAXN = origin.length + 1;
        arr = new int[MAXN];
        for (int i = 0; i < MAXN; i++) {
            arr[i] = origin[i - 1];
        }
        sum = new int[MAXN << 2];
        lazy = new int[MAXN << 2];
        change = new int[MAXN << 2];
        update = new boolean[MAXN << 2];
    }

    private void pushUp(int rt) {
        sum[rt] = sum[rt << 1] + sum[rt << 1 | 1];
    }

    /**
     * 下推
     * 
     * @param rt 当前的目标节点
     * @param ln 目标节点左子树的节点个数
     * @param rn 目标节点柚子树节点个数
     */
    private void pushDown(int rt, int ln, int rn) {
        if (update[rt]) {
            update[rt << 1] = true;
            update[rt << 1 | 1] = true;
            change[rt << 1] = change[rt];
            change[rt << 1 | 1] = change[rt];
            lazy[rt << 1] = 0;
            lazy[rt << 1 | 1] = 0;
            sum[rt << 1] = change[rt] * ln;
            sum[rt << 1 | 1] = change[rt] * rn;
            update[rt] = false;
        }
        // 下发上次更新后执行的add操作
        if (lazy[rt] != 0) {
            lazy[rt << 1] += lazy[rt];
            lazy[rt << 1 | 1] += lazy[rt];
            sum[rt << 1] += lazy[rt] * ln;
            sum[rt << 1 | 1] += lazy[rt] * rn;
            lazy[rt] = 0;
        }
    }

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

    public void update(int L, int R, int C, int l, int r, int rt) {
        // 全包了,只懒更新根节点
        if (L <= l && r <= R) {
            update[rt] = true;
            change[rt] = C;
            sum[rt] = C * (r - l + 1);
            lazy[rt] = 0;
            return;
        }
        // 需要下发
        int mid = (l + r) >> 1;
        pushDown(rt, mid - l + 1, r - mid);
        if (L <= mid) {
            update(L, R, C, l, mid, rt << 1);
        }
        if (R > mid) {
            update(L, R, C, mid + 1, r, rt << 1 | 1);
        }
    }

    /**
     * 懒添加
     * 
     * @param L  任务左
     * @param R  任务右
     * @param C  加的数值
     * @param l  收到任务的数组左
     * @param r  收到任务的数组右
     * @param rt 根节点
     */
    public void add(int L, int R, int C, int l, int r, int rt) {
        if (L <= l && r <= R) {
            // 任务将当前的数组全包了,不用下传了
            sum[rt] += C * (r - l + 1);
            lazy[rt] += C;
            return;
        }
        // 下发任务
        int mid = (l + r) >> 1;
        // 下发之前的lazy add任务
        pushDown(rt, mid - l + 1, r - mid);
        if (L <= mid) {
            add(L, R, C, l, mid, rt << 1);
        }
        if (R > mid) {
            add(L, R, C, mid + 1, r, rt << 1 | 1);
        }
        pushUp(rt);
    }

    public long query(int L, int R, int l, int r, int rt) {
        if (L <= l && r <= R) {
            return sum[rt];
        }
        int mid = (l + r) >> 1;
        pushDown(rt, mid - l + 1, r - mid);
        long ans = 0;
        if (L <= mid) {
            ans += query(L, R, l, mid, rt << 1);
        }
        if (R > mid) {
            ans += query(L, R, mid + 1, r, rt << 1 | 1);
        }
        return ans;
    }
}