sotanishy's competitive programming library

sotanishy's code snippets for competitive programming

View the Project on GitHub sotanishy/cp-library-cpp

:heavy_check_mark: test/yosupo/point_add_rectangle_sum.2d_segtree.test.cpp

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/point_add_rectangle_sum"

#include <bits/stdc++.h>

#include "../../data-structure/segtree/segment_tree_2d.hpp"
using namespace std;
using ll = long long;

struct AddMonoid {
    using T = ll;
    static T id() { return 0; }
    static T op(T a, T b) { return a + b; }
};

struct Query {
    int t, x, y, w, l, d, r, u;
    Query(int t, int x, int y, int w) : t(t), x(x), y(y), w(w) {}
    Query(int t, int l, int d, int r, int u) : t(t), l(l), d(d), r(r), u(u) {}
};

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    cin >> N >> Q;
    vector<int> x(N), y(N), w(N);
    vector<pair<int, int>> pts;
    for (int i = 0; i < N; ++i) {
        cin >> x[i] >> y[i] >> w[i];
        pts.push_back({x[i], y[i]});
    }
    vector<Query> queries;
    for (int i = 0; i < Q; ++i) {
        int t;
        cin >> t;
        if (t == 0) {
            int x, y, w;
            cin >> x >> y >> w;
            queries.emplace_back(t, x, y, w);
            pts.push_back({x, y});
        } else {
            int l, d, r, u;
            cin >> l >> d >> r >> u;
            queries.emplace_back(t, l, d, r, u);
        }
    }
    SegmentTree2D<int, int, AddMonoid> st(pts);
    for (int i = 0; i < N; ++i) {
        st.update(x[i], y[i], st.get(x[i], y[i]) + w[i]);
    }

    for (auto& q : queries) {
        if (q.t == 0) {
            st.update(q.x, q.y, st.get(q.x, q.y) + q.w);
        } else {
            cout << st.fold(q.l, q.r, q.d, q.u) << "\n";
        }
    }
}
#line 1 "test/yosupo/point_add_rectangle_sum.2d_segtree.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/point_add_rectangle_sum"

#include <bits/stdc++.h>

#line 3 "data-structure/segtree/segment_tree_2d.hpp"
#include <bit>
#line 7 "data-structure/segtree/segment_tree_2d.hpp"

#line 5 "data-structure/segtree/segment_tree.hpp"

template <typename M>
class SegmentTree {
    using T = M::T;

   public:
    SegmentTree() = default;
    explicit SegmentTree(int n) : SegmentTree(std::vector<T>(n, M::id())) {}
    explicit SegmentTree(const std::vector<T>& v)
        : size(std::bit_ceil(v.size())), node(2 * size, M::id()) {
        std::ranges::copy(v, node.begin() + size);
        for (int i = size - 1; i > 0; --i) {
            node[i] = M::op(node[2 * i], node[2 * i + 1]);
        }
    }

    T operator[](int k) const { return node[k + size]; }

    void update(int k, const T& x) {
        k += size;
        node[k] = x;
        while (k >>= 1) node[k] = M::op(node[2 * k], node[2 * k + 1]);
    }

    T fold(int l, int r) const {
        T vl = M::id(), vr = M::id();
        for (l += size, r += size; l < r; l >>= 1, r >>= 1) {
            if (l & 1) vl = M::op(vl, node[l++]);
            if (r & 1) vr = M::op(node[--r], vr);
        }
        return M::op(vl, vr);
    }

    template <typename F>
    int find_first(int l, F cond) const {
        T v = M::id();
        for (l += size; l > 0; l >>= 1) {
            if (l & 1) {
                T nv = M::op(v, node[l]);
                if (cond(nv)) {
                    while (l < size) {
                        nv = M::op(v, node[2 * l]);
                        if (cond(nv)) {
                            l = 2 * l;
                        } else {
                            v = nv, l = 2 * l + 1;
                        }
                    }
                    return l + 1 - size;
                }
                v = nv;
                ++l;
            }
        }
        return -1;
    }

    template <typename F>
    int find_last(int r, F cond) const {
        T v = M::id();
        for (r += size; r > 0; r >>= 1) {
            if (r & 1) {
                --r;
                T nv = M::op(node[r], v);
                if (cond(nv)) {
                    while (r < size) {
                        nv = M::op(node[2 * r + 1], v);
                        if (cond(nv)) {
                            r = 2 * r + 1;
                        } else {
                            v = nv, r = 2 * r;
                        }
                    }
                    return r - size;
                }
                v = nv;
            }
        }
        return -1;
    }

   private:
    int size;
    std::vector<T> node;
};
#line 9 "data-structure/segtree/segment_tree_2d.hpp"

template <typename X, typename Y, typename M>
class SegmentTree2D {
    using T = M::T;

   public:
    SegmentTree2D() = default;
    explicit SegmentTree2D(const std::vector<std::pair<X, Y>>& pts) {
        for (auto& [x, y] : pts) {
            xs.push_back(x);
        }
        std::ranges::sort(xs);
        xs.erase(std::ranges::unique(xs).begin(), xs.end());

        const int n = xs.size();
        size = std::bit_ceil((unsigned int)n);
        ys.resize(2 * size);
        seg.resize(2 * size);

        for (auto& [x, y] : pts) {
            ys[size + getx(x)].push_back(y);
        }

        for (int i = size + n - 1; i > 0; --i) {
            if (i >= size) {
                std::ranges::sort(ys[i]);
            } else {
                std::ranges::merge(ys[2 * i], ys[2 * i + 1],
                                   std::back_inserter(ys[i]));
            }
            ys[i].erase(std::ranges::unique(ys[i]).begin(), ys[i].end());
        }
        for (int i = 1; i < size + n; ++i) {
            seg[i] = SegmentTree<M>(ys[i].size());
        }
    }

    T get(X x, Y y) const {
        int kx = getx(x);
        assert(kx < (int)xs.size() && xs[kx] == x);
        kx += size;
        int ky = gety(kx, y);
        assert(ky < (int)ys[kx].size() && ys[kx][ky] == y);
        return seg[kx][ky];
    }

    void update(X x, Y y, T val) {
        int kx = getx(x);
        assert(kx < (int)xs.size() && xs[kx] == x);
        kx += size;
        int ky = gety(kx, y);
        assert(ky < (int)ys[kx].size() && ys[kx][ky] == y);
        seg[kx].update(ky, val);
        while (kx >>= 1) {
            ky = gety(kx, y);
            int kl = gety(2 * kx, y), kr = gety(2 * kx + 1, y);
            T vl = (kl < (int)ys[2 * kx].size() && ys[2 * kx][kl] == y
                        ? seg[2 * kx][kl]
                        : M::id());
            T vr = (kr < (int)ys[2 * kx + 1].size() && ys[2 * kx + 1][kr] == y
                        ? seg[2 * kx + 1][kr]
                        : M::id());
            seg[kx].update(ky, M::op(vl, vr));
        }
    }

    T fold(X sx, X tx, Y sy, Y ty) const {
        T ret = M::id();
        for (int l = size + getx(sx), r = size + getx(tx); l < r;
             l >>= 1, r >>= 1) {
            if (l & 1) {
                ret = M::op(ret, seg[l].fold(gety(l, sy), gety(l, ty)));
                ++l;
            }
            if (r & 1) {
                --r;
                ret = M::op(ret, seg[r].fold(gety(r, sy), gety(r, ty)));
            }
        }
        return ret;
    }

   private:
    int size;
    std::vector<X> xs;
    std::vector<std::vector<Y>> ys;
    std::vector<SegmentTree<M>> seg;

    int getx(X x) const { return std::ranges::lower_bound(xs, x) - xs.begin(); }
    int gety(int k, Y y) const {
        return std::ranges::lower_bound(ys[k], y) - ys[k].begin();
    }
};
#line 6 "test/yosupo/point_add_rectangle_sum.2d_segtree.test.cpp"
using namespace std;
using ll = long long;

struct AddMonoid {
    using T = ll;
    static T id() { return 0; }
    static T op(T a, T b) { return a + b; }
};

struct Query {
    int t, x, y, w, l, d, r, u;
    Query(int t, int x, int y, int w) : t(t), x(x), y(y), w(w) {}
    Query(int t, int l, int d, int r, int u) : t(t), l(l), d(d), r(r), u(u) {}
};

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    cin >> N >> Q;
    vector<int> x(N), y(N), w(N);
    vector<pair<int, int>> pts;
    for (int i = 0; i < N; ++i) {
        cin >> x[i] >> y[i] >> w[i];
        pts.push_back({x[i], y[i]});
    }
    vector<Query> queries;
    for (int i = 0; i < Q; ++i) {
        int t;
        cin >> t;
        if (t == 0) {
            int x, y, w;
            cin >> x >> y >> w;
            queries.emplace_back(t, x, y, w);
            pts.push_back({x, y});
        } else {
            int l, d, r, u;
            cin >> l >> d >> r >> u;
            queries.emplace_back(t, l, d, r, u);
        }
    }
    SegmentTree2D<int, int, AddMonoid> st(pts);
    for (int i = 0; i < N; ++i) {
        st.update(x[i], y[i], st.get(x[i], y[i]) + w[i]);
    }

    for (auto& q : queries) {
        if (q.t == 0) {
            st.update(q.x, q.y, st.get(q.x, q.y) + q.w);
        } else {
            cout << st.fold(q.l, q.r, q.d, q.u) << "\n";
        }
    }
}
Back to top page