lib/data_structure/dynamic_binary_indexed_tree_2d.hpp

• View this file on GitHub
• Last update: 2024-05-04 18:06:16+09:00
• Include: #include "lib/data_structure/dynamic_binary_indexed_tree_2d.hpp"

Depends on

• lib/data_structure/dynamic_binary_indexed_tree.hpp

Verified with

• test/library_checker/data_structure/point_add_rectangle_sum.test.cpp

Code

#pragma once

#include "../data_structure/dynamic_binary_indexed_tree.hpp"

/*
    DynamicBinaryIndexedTree2D<S, T>(h, w) : BIT2Dをサイズh*wで構築
    add(x, y, v) : (x, y) に v を加算します。
    sum(x1, y1, x2, y2) : [(x1, y1), (x2, y2)] の和を取得します。
*/

template <typename T>
struct DynamicBinaryIndexedTree2D{
    using BIT = DynamicBinaryIndexedTree<int, T>;
    int H, W;
    vector<BIT *> bit;
    DynamicBinaryIndexedTree2D(int h, int w) : H(h + 1), W(w){
        for(int i = 0; i < H; i++){
            bit.push_back(new BIT(W));
        }
    }

    void add(int i, int j, const T &x){
        for(++i; i < H; i += i & -i) (*bit[i]).add(j, x);
    }

    T sum(int x, int y){
        if(x < 0 || y < 0){
            return 0;
        }
        T res = 0;
        while(x > 0){
            res += (*bit[x]).sum(y);
            x -= x & -x;
        }
        return res;
    }

    T sum(int x1, int y1, int x2, int y2){
        T res = 0;
        while(x1 != x2){
            if(x1 < x2){
                res += (*bit[x2]).sum(y1, y2);
                x2 -= x2 & -x2;
            } else{
                res -= (*bit[x1]).sum(y1, y2);
                x1 -= x1 & -x1;
            }
        }
        return res;
    }
};

#line 2 "lib/data_structure/dynamic_binary_indexed_tree_2d.hpp"

#line 2 "lib/data_structure/dynamic_binary_indexed_tree.hpp"

/*
    DynamicBinaryIndexedTree<S, T>(n) : BITをサイズnで構築
    add(i, x) : i 番目の要素に加算します。 O(log(n)^2)
    sum(l, r) : [l, r) の区間和を取得します。O(log(n)^2)

    (以下の機能を使うには、0 <= i < N において、A_i >= 0 である必要があります。)
    lower_bound(x) : A_0 + A_1 + ... + A_y >= x となる最小の y を求めます。O(log(n))
*/

template <typename S, typename T>
struct DynamicBinaryIndexedTree{
    S N;
    unordered_map<S, T> data;
    DynamicBinaryIndexedTree(S _N) : N(_N + 1){
    }

    void add(S i, const T &x){
        for(++i; i < N; i += i & -i) data[i] += x;
    }

    T sum(int i){
        if(i < 0){
            return 0;
        }
        T ans = 0;
        while(i > 0){
            const auto iter = data.find(i);
            if(iter != data.end()){
                ans += iter->second;
            }
            i -= i & -i;
        }
        return ans;
    }

    T sum(int L, int R){
        return sum(R) - sum(L);
    }

    T operator[](S i){
        return sum(i + 1) - sum(i);
    }

    S lower_bound(T x){
        if(x <= 0){
            return 0;
        }
        S v = 0, r = 1;
        while(r < N) r = r << 1;
        for(S len = r; len > 0; len = len >> 1){
            if(v + len < N && data[v + len] < x){
                x -= data[v + len];
                v += len;
            }
        }
        return v;
    }
};
#line 4 "lib/data_structure/dynamic_binary_indexed_tree_2d.hpp"

/*
    DynamicBinaryIndexedTree2D<S, T>(h, w) : BIT2Dをサイズh*wで構築
    add(x, y, v) : (x, y) に v を加算します。
    sum(x1, y1, x2, y2) : [(x1, y1), (x2, y2)] の和を取得します。
*/

template <typename T>
struct DynamicBinaryIndexedTree2D{
    using BIT = DynamicBinaryIndexedTree<int, T>;
    int H, W;
    vector<BIT *> bit;
    DynamicBinaryIndexedTree2D(int h, int w) : H(h + 1), W(w){
        for(int i = 0; i < H; i++){
            bit.push_back(new BIT(W));
        }
    }

    void add(int i, int j, const T &x){
        for(++i; i < H; i += i & -i) (*bit[i]).add(j, x);
    }

    T sum(int x, int y){
        if(x < 0 || y < 0){
            return 0;
        }
        T res = 0;
        while(x > 0){
            res += (*bit[x]).sum(y);
            x -= x & -x;
        }
        return res;
    }

    T sum(int x1, int y1, int x2, int y2){
        T res = 0;
        while(x1 != x2){
            if(x1 < x2){
                res += (*bit[x2]).sum(y1, y2);
                x2 -= x2 & -x2;
            } else{
                res -= (*bit[x1]).sum(y1, y2);
                x1 -= x1 & -x1;
            }
        }
        return res;
    }
};

Back to top page