资瓷区间修改+区间求和的树状数组（一维/二维）

Posted 2020-12-31 arg-53

一维：令 (v_i) 为差分数组，那么 ([0, k]) 的前缀和就是 (sum{v_i(k+1-i)} = (k+1) cdot sum{v_i} + sum{v_i cdot (-i)})，树状数组维护一下 (v_i) 和 (v_i cdot i) 即可。

template <typename I>
struct Fenwick {
    struct Node {
        I v0, v1;

        void add(I d0, I d1) {
            v0 += d0;
            v1 += d1;
        }

        void operator += (const Node &rhs) {
            v0 += rhs.v0;
            v1 += rhs.v1;
        }
    };

    Fenwick(int n) : n(n), tree(n) {}

    void orz(int k, I d0, I d1) {
        for (; k < n; k |= k + 1) {
            tree[k].add(d0, d1);
        }
    }

    void add(int l, int r, I d) { // [l, r)
        orz(l, d, -l * d);
        orz(r, -d, r * d);
    }

    I sum(int k) { // [0, k]
        Node res = {};
        for (int i = k; i >= 0; i = (i & (i + 1)) - 1) {
            res += tree[i];     
        }
        return (k + 1) * res.v0 + res.v1;
    }

    I sum(int l, int r) { // [l, r)
        return sum(r - 1) - sum(l - 1);
    }

    int n;
    vector<Node> tree;
};

二维：和一维的推导类似，维护一下 (v(i,j))，(v(i,j) cdot i)，(v(i,j) cdot j) 和 (v(i,j) cdot ij)。

struct Fenwick {
    struct Node {
        int v, vi, vj, vij;

        void operator += (const Node &rhs) {
            v += rhs.v;
            vi += rhs.vi;
            vj += rhs.vj;
            vij += rhs.vij;
        }

        void apply(int d, int di, int dj, int dij) {
            v += d;
            vi += di;
            vj += dj;
            vij += dij;
        }
    };

    Fenwick(int n, int m) : n(n), m(m), tree(n, vector<Node>(m)) {}

    void add(int x, int y, int d) {
        int di = -x * d;
        int dj = -y * d;
        int dij = x * y * d;
        for (int i = x; i < n; i |= i + 1) {
            for (int j = y; j < m; j |= j + 1) {
                tree[i][j].apply(d, di, dj, dij);
            }
        }
    }

    int sum(int x, int y) {
        Node res = {};   
        for (int i = x; i >= 0; i = (i & (i + 1)) - 1) {
            for (int j = y; j >= 0; j = (j & (j + 1)) - 1) {
                res += tree[i][j];
            }
        }
        return (x + 1) * (y + 1) * res.v + (y + 1) * res.vi + (x + 1) * res.vj + res.vij;
    }
    
    int n, m;
    vector< vector<Node> > tree;
};