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/rectangle_sum.persistent_segment_tree.test.cpp

Depends on

Code

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

#include <bits/stdc++.h>

#include "../../data-structure/segtree/persistent_segment_tree.hpp"
#include "../../misc/compress.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; }
};

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

    int N, Q;
    cin >> N >> Q;
    vector<int> xs, ys;
    vector<int> x(N), y(N), w(N);
    vector<int> l(Q), d(Q), r(Q), u(Q);
    for (int i = 0; i < N; ++i) {
        cin >> x[i] >> y[i] >> w[i];
        xs.push_back(x[i]);
        ys.push_back(y[i]);
    }
    for (int i = 0; i < Q; ++i) {
        cin >> l[i] >> d[i] >> r[i] >> u[i];
        xs.push_back(l[i]);
        ys.push_back(d[i]);
        xs.push_back(r[i]);
        ys.push_back(u[i]);
    }
    xs.push_back(-1);
    ys.push_back(-1);
    Compress<int> cx(xs), cy(ys);
    x = cx.compress(x);
    y = cy.compress(y);
    l = cx.compress(l);
    d = cy.compress(d);
    r = cx.compress(r);
    u = cy.compress(u);
    int M = xs.size();
    vector<vector<pair<int, int>>> points(M);
    for (int i = 0; i < N; ++i) {
        points[x[i]].emplace_back(y[i], w[i]);
    }
    vector<PersistentSegmentTree<AddMonoid>> st(M);
    st[0] = PersistentSegmentTree<AddMonoid>(M);
    for (int x = 0; x < M; ++x) {
        if (x > 0) st[x] = st[x - 1];
        auto tmp = st.back();
        for (auto& p : points[x]) {
            int y, w;
            tie(y, w) = p;
            st[x] = st[x].update(y, st[x][y] + w);
        }
    }
    for (int i = 0; i < Q; ++i) {
        ll ans = st[r[i] - 1].fold(d[i], u[i]) - st[l[i] - 1].fold(d[i], u[i]);
        cout << ans << "\n";
    }
}
#line 1 "test/yosupo/rectangle_sum.persistent_segment_tree.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/rectangle_sum"

#include <bits/stdc++.h>

#line 2 "data-structure/segtree/persistent_segment_tree.hpp"
#include <bit>
#line 6 "data-structure/segtree/persistent_segment_tree.hpp"

/**
 * @brief Persistent Segment Tree
 */
template <typename M>
class PersistentSegmentTree {
    using T = M::T;

   public:
    PersistentSegmentTree() = default;
    explicit PersistentSegmentTree(int n)
        : PersistentSegmentTree(std::vector<T>(n, M::id())) {}
    explicit PersistentSegmentTree(const std::vector<T>& v)
        : root(std::make_shared<Node>()), size(std::bit_ceil(v.size())) {
        build(v, root, 0, size);
    }

    T operator[](int k) const { return fold(k, k + 1); }

    PersistentSegmentTree update(int k, const T& x) const {
        return PersistentSegmentTree(update(k, x, root, 0, size), size);
    }

    T fold(int l, int r) const { return fold(l, r, root, 0, size); }

   private:
    struct Node;
    using node_ptr = std::shared_ptr<Node>;

    struct Node {
        T val;
        node_ptr left, right;
        Node() : val(M::id()), left(nullptr), right(nullptr) {}
        Node(const T& val, const node_ptr& left, const node_ptr& right)
            : val(val), left(left), right(right) {}
    };

    node_ptr root;
    int size;

    PersistentSegmentTree(const node_ptr& root, int size)
        : root(root), size(size) {}

    void build(const std::vector<T>& v, const node_ptr& n, int l, int r) const {
        if (r - l == 1) {
            n->val = l < (int)v.size() ? v[l] : M::id();
            return;
        }
        int m = std::midpoint(l, r);
        n->left = std::make_shared<Node>();
        build(v, n->left, l, m);
        n->right = std::make_shared<Node>();
        build(v, n->right, m, r);
        n->val = M::op(n->left->val, n->right->val);
    }

    node_ptr update(int k, const T& x, const node_ptr& n, int l, int r) const {
        if (r - l == 1) {
            return std::make_shared<Node>(x, nullptr, nullptr);
        }
        int m = std::midpoint(l, r);
        if (k < m) {
            auto left = update(k, x, n->left, l, m);
            T val = M::op(left->val, n->right->val);
            return std::make_shared<Node>(val, left, n->right);
        } else {
            auto right = update(k, x, n->right, m, r);
            T val = M::op(n->left->val, right->val);
            return std::make_shared<Node>(val, n->left, right);
        }
    }

    T fold(int a, int b, const node_ptr& n, int l, int r) const {
        if (r <= a || b <= l) return M::id();
        if (a <= l && r <= b) return n->val;
        int m = std::midpoint(l, r);
        return M::op(fold(a, b, n->left, l, m), fold(a, b, n->right, m, r));
    }
};
#line 4 "misc/compress.hpp"

/**
 * @brief Coordinate Compression
 */
template <typename T>
class Compress {
   public:
    Compress() = default;
    explicit Compress(const std::vector<T>& vs) : xs(vs) {
        std::ranges::sort(xs);
        xs.erase(std::ranges::unique(xs).begin(), xs.end());
    }

    int compress(const T& x) const {
        return std::ranges::lower_bound(xs, x) - xs.begin();
    }

    std::vector<int> compress(const std::vector<T>& vs) const {
        std::vector<int> res(vs.size());
        std::ranges::transform(vs, res.begin(),
                               [&](const T& x) { return compress(x); });
        return res;
    }

    T decompress(int i) const { return xs[i]; }

    int size() const { return xs.size(); }

   private:
    std::vector<T> xs;
};
#line 7 "test/yosupo/rectangle_sum.persistent_segment_tree.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; }
};

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

    int N, Q;
    cin >> N >> Q;
    vector<int> xs, ys;
    vector<int> x(N), y(N), w(N);
    vector<int> l(Q), d(Q), r(Q), u(Q);
    for (int i = 0; i < N; ++i) {
        cin >> x[i] >> y[i] >> w[i];
        xs.push_back(x[i]);
        ys.push_back(y[i]);
    }
    for (int i = 0; i < Q; ++i) {
        cin >> l[i] >> d[i] >> r[i] >> u[i];
        xs.push_back(l[i]);
        ys.push_back(d[i]);
        xs.push_back(r[i]);
        ys.push_back(u[i]);
    }
    xs.push_back(-1);
    ys.push_back(-1);
    Compress<int> cx(xs), cy(ys);
    x = cx.compress(x);
    y = cy.compress(y);
    l = cx.compress(l);
    d = cy.compress(d);
    r = cx.compress(r);
    u = cy.compress(u);
    int M = xs.size();
    vector<vector<pair<int, int>>> points(M);
    for (int i = 0; i < N; ++i) {
        points[x[i]].emplace_back(y[i], w[i]);
    }
    vector<PersistentSegmentTree<AddMonoid>> st(M);
    st[0] = PersistentSegmentTree<AddMonoid>(M);
    for (int x = 0; x < M; ++x) {
        if (x > 0) st[x] = st[x - 1];
        auto tmp = st.back();
        for (auto& p : points[x]) {
            int y, w;
            tie(y, w) = p;
            st[x] = st[x].update(y, st[x][y] + w);
        }
    }
    for (int i = 0; i < Q; ++i) {
        ll ans = st[r[i] - 1].fold(d[i], u[i]) - st[l[i] - 1].fold(d[i], u[i]);
        cout << ans << "\n";
    }
}
Back to top page