526互联

【个人模板封装】树套树、高维数据结构

发布时间 2023-07-30 21:19:57作者: hh2048

前言

这是我个人使用的一些模板封装，限于个人能力，可能存在诸多不足与漏洞，在未加测试直接使用前请务必小心谨慎。更新可能会滞后于我本地的文档，如有疑问或者催更之类的可以在评论区留言。

全文模板测试均基于以下版本信息，请留意版本兼容问题。

Windows, 64bit
G++ (ISO C++20)
stack=268435456
开启O2优化

树状数组部分

树状数组

struct BIT {
    int n;
    vector<int> w;
    
    BIT() {}
    BIT(int n) {
        this->n = n; // 这里必须写 n ，否则会RE
        w.resize(n + 1);
    }
    void add(int x, int k) {
        for (; x <= n; x += x & -x) {
            w[x] += k;
        }
    }
    void add(int x, int y, int k) {
        add(x, k), add(y, -k);
    }
    int ask(int x) {
        int ans = 0;
        for (; x; x -= x & -x) {
            ans += w[x];
        }
        return ans;
    }
    int ask(int x, int y) {
        return ask(y) - ask(x - 1);
    }
    int kth(int k) { // 查找第k大的值
        int ans = 0;
        for (int i = __lg(n); i >= 0; i--) {
            int val = ans + (1 << i);
            if (val < n && w[val] < k) {
                k -= w[val];
                ans = val;
            }
        }
        return ans + 1;
    }
};

树状数组套树状数组（二维树状数组）1

请注意，该版本不能同时进行区间修改+区间查询。无离散化版本的空间占用为 \(\mathcal O(NM)\) 、建树复杂度为 \(\mathcal O(NM)\) 、单次查询复杂度为 \(\mathcal O(\log N\cdot \log M)\) 。

大致有以下两种写法，均通过模板题测试。

该部分模板题可参考：LOJ #133. 二维树状数组 1：单点修改，区间查询、LOJ #134. 二维树状数组 2：区间修改，单点查询。

借助一维树状数组进行拓展

struct BIT_2D {
    struct BIT {
        int n;
        vector<int> w;
        BIT() {}
        BIT(int n) {
            this->n = n; // 这里必须写 n ，否则会RE
            w.resize(n + 1);
        }
        void modify(int x, int k) {
            for (; x <= n; x += x & -x) {
                w[x] += k;
            }
        }
        int ask(int x) {
            int ans = 0;
            for (; x; x -= x & -x) {
                ans += w[x];
            }
            return ans;
        }
    };
    int n;
    vector<BIT> w;
    
    BIT_2D() {}
    BIT_2D(int n, int m) {
        this->n = n;
        w.resize(n + 1);
        for (int i = 1; i <= n; i++) {
            w[i] = BIT(m);
        }
    }
    void modify(int x, int y, int k) {
        for (; x <= n; x += x & -x) {
            w[x].modify(y, k);
        }
    }
    void modify(int x, int y, int X, int Y, int k) { // 区块修改：二维差分
        X++, Y++;
        modify(x, y, k), modify(X, y, -k);
        modify(X, Y, k), modify(x, Y, -k);
    }
    int ask(int x, int y) { // 单点查询
        int ans = 0;
        for (; x; x -= x & -x) {
            ans += w[x].ask(y);
        }
        return ans;
    }
    int ask(int x, int y, int X, int Y) { // 区块查询：二维前缀和
        x--, y--;
        return ask(X, Y) - ask(x, Y) - ask(X, y) + ask(x, y);
    }
};

压缩优化

struct BIT_2D {
    int n, m;
    vector<vector<int>> w;
    
    BIT_2D(int n, int m) : n(n), m(m) {
        w.resize(n + 1, vector<int>(m + 1));
    }
    void add(int x, int y, int k) {
        for (int i = x; i <= n; i += i & -i) {
            for (int j = y; j <= m; j += j & -j) {
                w[i][j] += k;
            }
        }
    }
    void add(int x, int y, int X, int Y, int k) { // 区块修改：二维差分
        X++, Y++;
        add(x, y, k), add(X, y, -k);
        add(X, Y, k), add(x, Y, -k);
    }
    int ask(int x, int y) { // 单点查询
        int ans = 0;
        for (int i = x; i; i -= i & -i) {
            for (int j = y; j; j -= j & -j) {
                ans += w[i][j];
            }
        }
        return ans;
    }
    int ask(int x, int y, int X, int Y) { // 区块查询：二维前缀和
        x--, y--;
        return ask(X, Y) - ask(x, Y) - ask(X, y) + ask(x, y);
    }
};

树状数组套树状数组（二维树状数组）2

该版本支持全部操作。但是全部复杂度均比上一个版本多 \(4\) 倍。

这里仅提供压缩优化版。

该部分模板题可参考：LOJ #135. 二维树状数组 3：区间修改，区间查询。

struct BIT_2D {
    int n, m;
    vector<vector<int>> b1, b2, b3, b4;
    
    BIT_2D(int n, int m) : n(n), m(m) {
        b1.resize(n + 1, vector<int>(m + 1));
        b2.resize(n + 1, vector<int>(m + 1));
        b3.resize(n + 1, vector<int>(m + 1));
        b4.resize(n + 1, vector<int>(m + 1));
    }
    void add(auto &w, int x, int y, int k) { // 单点修改
        for (int i = x; i <= n; i += i & -i) {
            for (int j = y; j <= m; j += j & -j) {
                w[i][j] += k;
            }
        }
    }
    void add(int x, int y, int k) { // 多了一步计算
        add(b1, x, y, k);
        add(b2, x, y, k * (x - 1));
        add(b3, x, y, k * (y - 1));
        add(b4, x, y, k * (x - 1) * (y - 1));
    }
    void add(int x, int y, int X, int Y, int k) { // 区块修改：二维差分
        X++, Y++;
        add(x, y, k), add(X, y, -k);
        add(X, Y, k), add(x, Y, -k);
    }
    int ask(auto &w, int x, int y) { // 单点查询
        int ans = 0;
        for (int i = x; i; i -= i & -i) {
            for (int j = y; j; j -= j & -j) {
                ans += w[i][j];
            }
        }
        return ans;
    }
    int ask(int x, int y) { // 多了一步计算
        int ans = 0;
        ans += x * y * ask(b1, x, y);
        ans -= y * ask(b2, x, y);
        ans -= x * ask(b3, x, y);
        ans += ask(b4, x, y);
        return ans;
    }
    int ask(int x, int y, int X, int Y) { // 区块查询：二维前缀和
        x--, y--;
        return ask(X, Y) - ask(x, Y) - ask(X, y) + ask(x, y);
    }
};