src/segtree2d.hpp

• View this file on GitHub
• Last update: 2025-10-08 01:18:35+09:00
• Include: #include "src/segtree2d.hpp"

Depends on

• Segment Tree (src/segtree.hpp)
• src/template.hpp

Verified with

• src/test/yosupo_point_add_rectangle_sum.cpp

Code

#include "segtree.hpp"

template<typename K, class S, S (*op)(S, S), S (*e)()> struct segtree_2d {
    int n;
    vector<segtree<S, op, e>> d;
    vector<vector<K>> ys;
    vector<pair<K, K>> ps;
    segtree_2d() {}
    void add_point(K x, K y) { ps.push_back({x, y}); }

    void build() {
        sort(all(ps));
        ps.erase(unique(all(ps)), ps.end());
        n = __bit_ceil(si(ps));
        ys.resize(2 * n);
        rep(i, si(ps)) ys[n + i] = {ps[i].se};
        for(int i = n; --i;) merge(all(ys[i * 2]), all(ys[i * 2 + 1]), back_inserter(ys[i]));

        d.resize(2 * n);
        for(int i = 1; i < 2 * n; i++) {
            d[i] = segtree<S, op, e>(si(ys[i]) + 1);
        }
    }

    void set(K x, K y, S v) {
        auto it = lower_bound(all(ps), make_pair(x, y));
        assert(it != ps.end() && it->fi == x && it->se == y);
        int k = it - ps.begin();
        for(k += n; k; k >>= 1) {
            auto it = lower_bound(all(ys[k]), y);
            assert(it != ys[k].end());
            int id = it - ys[k].begin();
            d[k].set(id, v);
        }
    }
    void add(K x, K y, S v) {
        auto it = lower_bound(all(ps), make_pair(x, y));
        assert(it != ps.end() && it->fi == x && it->se == y);
        int k = it - ps.begin();
        for(k += n; k; k >>= 1) {
            auto it = lower_bound(all(ys[k]), y);
            assert(it != ys[k].end());
            int id = it - ys[k].begin();
            d[k].set(id, op(d[k].get(id), v));
        }
    }

    S range(K lv, K rv, K dv, K uv) {
        if(dv >= uv) return e();
        int l = lower_bound(all(ps), make_pair(lv, -1)) - ps.begin();
        int r = lower_bound(all(ps), make_pair(rv, -1)) - ps.begin();

        S sml = e(), smr = e();
        for(l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if(l & 1) {
                int li = lower_bound(all(ys[l]), dv) - ys[l].begin();
                int ri = lower_bound(all(ys[l]), uv) - ys[l].begin();
                sml = op(sml, d[l].prod(li, ri));
                l++;
            }
            if(r & 1) {
                --r;
                int li = lower_bound(all(ys[r]), dv) - ys[r].begin();
                int ri = lower_bound(all(ys[r]), uv) - ys[r].begin();
                smr = op(d[r].prod(li, ri), smr);
            }
        }
        return op(sml, smr);
    }
};

#line 2 "src/template.hpp"
#include <bits/stdc++.h>
using namespace std;
#define si(a) (long)a.size()
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define rep(i, n) for(long i = 0; i < (long)(n); ++i)
template<typename T> bool chmin(T& a, T b) { return b < a ? (a = b, 1) : 0; }
template<typename T> bool chmax(T& a, T b) { return b > a ? (a = b, 1) : 0; }
struct _ {
    _() { cin.tie(0)->sync_with_stdio(0), cout.tie(0), cout << fixed << setprecision(16); }
} __;
#line 3 "src/segtree.hpp"

template<class S, S (*op)(S, S), S (*e)()> struct segtree {
    int n;
    vector<S> d;
    segtree(): segtree(1) {}
    explicit segtree(int n_): segtree(vector<S>(n_, e())) {}
    explicit segtree(const vector<S>& a): n(__bit_ceil(si(a))), d(n * 2, e()) {
        rep(i, si(a)) d[n + i] = a[i];
        for(int i = n; --i;) d[i] = op(d[i * 2], d[i * 2 + 1]);
    }
    void set(int i, S x) {
        for(d[i += n] = x; i >>= 1;) d[i] = op(d[i * 2], d[i * 2 + 1]);
    }
    S get(int i) { return d[i + n]; }
    S prod(int l, int r) {
        S L = e(), R = e();
        for(l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if(l & 1) L = op(L, d[l++]);
            if(r & 1) R = op(d[--r], R);
        }
        return op(L, R);
    }
    template<typename F> int max_right(int l, F f) const {
        if(l == n) return n;
        l += n;
        S sm = e();
        do {
            while(~l & 1) l >>= 1;
            if(!f(op(sm, d[l]))) {
                while(l < n) {
                    l <<= 1;
                    if(f(op(sm, d[l]))) sm = op(sm, d[l++]);
                }
                return l - n;
            }
            sm = op(sm, d[l++]);
        } while((l & -l) != l);
        return n;
    }
    template<typename F> int min_left(int r, F f) const {
        if(!r) return 0;
        r += n;
        S sm = e();
        do {
            r--;
            while(r > 1 and r & 1) r >>= 1;
            if(!f(op(d[r], sm))) {
                while(r < n) {
                    r = (2 * r + 1);
                    if(f(op(d[r], sm))) sm = op(d[r--], sm);
                }
                return r + 1 - n;
            }
            sm = op(d[r], sm);
        } while((r & -r) != r);
        return 0;
    }
};

template<class T> using SegtreeFrom = segtree<typename T::S, T::op, T::e>;
#line 2 "src/segtree2d.hpp"

template<typename K, class S, S (*op)(S, S), S (*e)()> struct segtree_2d {
    int n;
    vector<segtree<S, op, e>> d;
    vector<vector<K>> ys;
    vector<pair<K, K>> ps;
    segtree_2d() {}
    void add_point(K x, K y) { ps.push_back({x, y}); }

    void build() {
        sort(all(ps));
        ps.erase(unique(all(ps)), ps.end());
        n = __bit_ceil(si(ps));
        ys.resize(2 * n);
        rep(i, si(ps)) ys[n + i] = {ps[i].se};
        for(int i = n; --i;) merge(all(ys[i * 2]), all(ys[i * 2 + 1]), back_inserter(ys[i]));

        d.resize(2 * n);
        for(int i = 1; i < 2 * n; i++) {
            d[i] = segtree<S, op, e>(si(ys[i]) + 1);
        }
    }

    void set(K x, K y, S v) {
        auto it = lower_bound(all(ps), make_pair(x, y));
        assert(it != ps.end() && it->fi == x && it->se == y);
        int k = it - ps.begin();
        for(k += n; k; k >>= 1) {
            auto it = lower_bound(all(ys[k]), y);
            assert(it != ys[k].end());
            int id = it - ys[k].begin();
            d[k].set(id, v);
        }
    }
    void add(K x, K y, S v) {
        auto it = lower_bound(all(ps), make_pair(x, y));
        assert(it != ps.end() && it->fi == x && it->se == y);
        int k = it - ps.begin();
        for(k += n; k; k >>= 1) {
            auto it = lower_bound(all(ys[k]), y);
            assert(it != ys[k].end());
            int id = it - ys[k].begin();
            d[k].set(id, op(d[k].get(id), v));
        }
    }

    S range(K lv, K rv, K dv, K uv) {
        if(dv >= uv) return e();
        int l = lower_bound(all(ps), make_pair(lv, -1)) - ps.begin();
        int r = lower_bound(all(ps), make_pair(rv, -1)) - ps.begin();

        S sml = e(), smr = e();
        for(l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if(l & 1) {
                int li = lower_bound(all(ys[l]), dv) - ys[l].begin();
                int ri = lower_bound(all(ys[l]), uv) - ys[l].begin();
                sml = op(sml, d[l].prod(li, ri));
                l++;
            }
            if(r & 1) {
                --r;
                int li = lower_bound(all(ys[r]), dv) - ys[r].begin();
                int ri = lower_bound(all(ys[r]), uv) - ys[r].begin();
                smr = op(d[r].prod(li, ri), smr);
            }
        }
        return op(sml, smr);
    }
};

Back to top page