test/library-checker/DataStructure/RectangleSum_2.test.cpp

• View this file on GitHub
• Last update: 2025-11-24 18:49:07+09:00
• Problem: https://judge.yosupo.jp/problem/rectangle_sum

Depends on

• build/pch/stdc++.hpp
• library/datastructure/CumulativeSum.hpp
• library/datastructure/FullyIndexableDictionary.hpp
• library/datastructure/WaveletMatrix.hpp
• library/r2/Projection.hpp
• library/r2/XY.hpp
• library/util/Compress.hpp

Code

#define PROBLEM "https://judge.yosupo.jp/problem/rectangle_sum"
#include <bits/stdc++.h>

#define REP(i, n) for (int i = 0; i < (n); i++)

#include "library/datastructure/CumulativeSum.hpp"
#include "library/datastructure/WaveletMatrix.hpp"
#include "library/r2/Projection.hpp"
#include "library/r2/XY.hpp"

using ll = long long;
using r2 = XY<ll>;

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

    int n, q;
    std::cin >> n >> q;
    std::vector<r2> r2s(n);
    std::vector<int> x(n), y(n), w(n);
    REP (i, n) {
        std::cin >> x[i] >> y[i] >> w[i];
        r2s[i] = r2(x[i], y[i]);
    }

    auto P = Projection(r2s);
    std::vector<int> Y(P.size(), 0);
    REP (id, P.size())
        Y[id] = P.r(id).y;
    WaveletMatrix WM(Y);

    int LOG = WM.height();
    std::vector W(LOG, std::vector<ll>(P.size(), 0));
    REP (i, n) {
        int idx = P.id(x[i], y[i]);
        auto idxs = WM.points(idx);
        REP (h, LOG)
            W[h][idxs[h]] += w[i];
    }
    std::vector<CumulativeSum<ll>> C(LOG);
    REP (h, LOG)
        C[h] = CumulativeSum<ll>(W[h]);

    REP (j, q) {
        int l, r, d, u;
        std::cin >> l >> d >> r >> u;
        auto [L, R] = P.interval(l, r);
        ll res = 0;
        auto intervals = WM.intervals(L, R, u);
        for (const auto &[h, l, r] : intervals)
            res += C[h].sum(l, r);
        auto intervals2 = WM.intervals(L, R, d);
        for (const auto &[h, l, r] : intervals2)
            res -= C[h].sum(l, r);
        std::cout << res << "\n";
    }
}

#line 1 "test/library-checker/DataStructure/RectangleSum_2.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/rectangle_sum"
#include <bits/stdc++.h>

#define REP(i, n) for (int i = 0; i < (n); i++)

#line 1 "library/datastructure/CumulativeSum.hpp"
template <typename T> struct CumulativeSum {
    using U = std::conditional_t<std::is_same_v<T, int>, long long, T>;
    std::vector<U> A;
    CumulativeSum() : A(1, 0) {}
    CumulativeSum(const std::vector<T> &v) : A(v.size() + 1, 0) {
        for (int i = 0; i < v.size(); i++)
            A[i + 1] = A[i] + v[i];
    }
    void add(const T &a) { A.push_back(A.back() + a); }
    U sum(int l, int r) {
        assert(0 <= l and l <= r and r < A.size());
        return A[r] - A[l];
    }
    U sum() { return A.back(); }
};
#line 4 "library/datastructure/FullyIndexableDictionary.hpp"
class FullyIndexableDictionary {
    int n,
        block; // 64個事に区切ったブロックの個数
    std::vector<unsigned long long> bit;
    std::vector<unsigned int> sum; // ブロック毎の累積和
    bool prepared;

  public:
    FullyIndexableDictionary() {}
    FullyIndexableDictionary(int n)
        : n(n), block((n + 63) >> 6), bit(block, 0), sum(block + 1, 0),
          prepared(false) {}

    bool is_prepared() { return prepared; }

    void set(int k) {
        bit[k >> 6] |= 1ull << (k & 63);
        sum[(k >> 6) + 1]++;
    }
    void build() {
        assert(!prepared);
        prepared = true;
        for (int i = 0; i < block; i++)
            sum[i + 1] += sum[i];
    }

    bool operator[](int k) const { return bool((bit[k >> 6] >> (k & 63)) & 1); }

    // [0,j) の合計
    int rank(int j, bool f = 1) {
        assert(prepared);
        int a = sum[j >> 6] +
                __builtin_popcountll(bit[j >> 6] & ((1ull << (j & 63)) - 1));
        return (f ? a : j - a);
    }
    // 0-indexed で k 番目の f の場所
    int select(int k, bool f = 1) {
        assert(prepared);
        if (k < 0 or rank(n, f) <= k)
            return -1;
        int l = 0, r = n;
        while (r - l > 1) {
            int m = (l + r) >> 1;
            (rank(m, f) >= k + 1 ? r : l) = m;
        }
        return r - 1;
    }
    // l以上で k 番目の f の場所
    int select(int k, bool f, int l) { return select(rank(l, f) + k, f); }
};
#line 2 "library/util/Compress.hpp"
template <typename T, bool Sentinel = false> class Compress {
    std::vector<T> v;
    bool prepared;

  public:
    Compress() : prepared(false) {
        if constexpr (Sentinel) {
            static_assert(std::numeric_limits<T>::is_specialized,
                          "cannot use Sentinel");
            v = {std::numeric_limits<T>::min(), std::numeric_limits<T>::max()};
        }
    }
    Compress(const std::vector<T> &w) : v(w), prepared(false) {
        if constexpr (Sentinel) {
            static_assert(std::numeric_limits<T>::is_specialized,
                          "cannot use Sentinel");
            v.push_back(std::numeric_limits<T>::min());
            v.push_back(std::numeric_limits<T>::max());
        }
        build();
    }

    void add(T a) {
        assert(!prepared);
        v.push_back(a);
    }
    void build() {
        assert(!prepared);
        prepared = true;
        std::ranges::sort(v);
        auto result = std::ranges::unique(v);
        v.erase(result.begin(), result.end());
    }

    bool is_prepared() const { return prepared; }

    int operator[](const T &a) const {
        assert(prepared);
        auto it = std::ranges::lower_bound(v, a);
        assert(*it == a);
        return std::distance(v.begin(), it);
    }
    int geq(const T &a) const {
        assert(prepared);
        auto it = std::ranges::lower_bound(v, a);
        return std::distance(v.begin(), it);
    }
    int gt(const T &a) const {
        assert(prepared);
        auto it = std::ranges::upper_bound(v, a);
        return std::distance(v.begin(), it);
    }
    int leq(const T &a) const {
        assert(prepared);
        auto it = --std::ranges::upper_bound(v, a);
        return std::distance(v.begin(), it);
    }
    int lt(const T &a) const {
        assert(prepared);
        auto it = --std::ranges::lower_bound(v, a);
        return std::distance(v.begin(), it);
    }
    T r(int id) const {
        assert(prepared);
        return v[id];
    }
    bool exist(const T &a) const {
        assert(prepared);
        return (*std::ranges::lower_bound(v, a)) == a;
    }
    int size() const { return v.size(); }
    T max() const { return v.back(); }
    T min() const { return v[0]; }

    friend std::ostream &operator<<(std::ostream &os, const Compress &C) {
        for (int i = 0; i < C.v.size(); i++)
            os << C.v[i] << ":" << i << " ";
        return os;
    }
};
#line 4 "library/datastructure/WaveletMatrix.hpp"
#define REP_(i, n) for (int i = 0; i < (n); i++)
template <typename T, bool COMPRESS = true> class WaveletMatrix {
  protected:
    using U = std::conditional_t<COMPRESS, int, T>;
    static_assert(std::is_integral_v<U>, "Wavelet Matrix is only for integer");
    int n, memo, log;
    std::vector<FullyIndexableDictionary> mat;
    std::vector<int> zero_cnt;
    Compress<T, true> C;
    std::vector<T> data;

    constexpr U comp(const T &x) const {
        if constexpr (COMPRESS) {
            return C.geq(x);
        } else {
            return x;
        }
    }
    constexpr T uncomp(const U &a) {
        if constexpr (COMPRESS) {
            return C.r(a);
        } else {
            return a;
        }
    }

    // 0-indexed で下から i bit 目
    inline bool low_bit(const U &a, int i) const { return (a >> i) & 1; }
    // 0-indexed で上から i bit 目
    inline bool high_bit(const U &a, int i) const {
        return low_bit(a, log - i - 1);
    }

    int nxt(int idx, int h, const U &a) {
        // idx の位置に a があった場合上から h bit 目でどこに行くか
        bool bit = high_bit(a, h);
        return mat[h].rank(idx, bit) + (bit ? zero_cnt[h] : 0);
    }

  public:
    WaveletMatrix(std::vector<T> v, int log_ = 0)
        : n(v.size()), data(v), log(log_) {
        std::vector<U> cv(n);
        if constexpr (COMPRESS) {
            C = Compress<T, true>(v);
            for (int i = 0; i < n; i++)
                cv[i] = C[v[i]];
            while (C.size() >= (1ull << log))
                log++;
        } else {
            cv = v;
            T mx = 0;
            for (const T &a : v) {
                assert(a >= 0);
                if (mx < a)
                    mx = a;
            }
            while (mx >= (1ull << log))
                log++;
        }
        mat.resize(log);
        zero_cnt.resize(log);
        std::vector<U> lv(n), rv(n);
        REP_(h, log) {
            mat[h] = FullyIndexableDictionary(n);
            int l = 0, r = 0;
            REP_(i, n)
            if (high_bit(cv[i], h)) {
                rv[r++] = cv[i];
                mat[h].set(i);
            } else
                lv[l++] = cv[i];
            zero_cnt[h] = l;
            mat[h].build();
            std::swap(lv, cv);
            REP_(i, r) cv[l + i] = rv[i];
        }
    }

    // [l,r) の x の個数
    int rank(int l, int r, const T &x) {
        if constexpr (COMPRESS) {
            if (!C.exist(x))
                return 0;
        }
        U a = comp(x);
        REP_(h, log) {
            l = nxt(l, h, a);
            r = nxt(r, h, a);
        }
        memo = l;
        return r - l;
    }
    int rank(int r, const T &x) { return rank(x, 0, r); }

    // k 番目の x の index
    int select(int k, const T &x) {
        U a = comp(x);
        if (rank(a, n) < k)
            return -1;
        k += memo;
        for (int h = log - 1; h >= 0; h--) {
            bool bit = high_bit(a, h);
            if (bit)
                k -= zero_cnt[h];
            k = mat[h].select(k, bit);
        }
        return k;
    }

    // [l,r) で 0-indexed k 番目に大きい値
    T kth_largest(int l, int r, int k) {
        if (k < 0 or r - l <= k)
            return -1;
        U res = 0;
        REP_(h, log) {
            int L = mat[h].rank(l);
            int R = mat[h].rank(r);
            res <<= 1;
            if (R - L > k) {
                l = L + zero_cnt[h];
                r = R + zero_cnt[h];
                res++;
            } else {
                k -= R - L;
                l -= L;
                r -= R;
            }
        }
        return uncomp(res);
    }
    T kth_smallest(int l, int r, int k) {
        return kth_largest(l, r, r - l - k - 1);
    }

    // [l,r) で x 未満の個数
    int range_freq(int l, int r, const T &upper) {
        U a = comp(upper);
        int res = 0;
        REP_(h, log) {
            if (high_bit(a, h)) {
                int L = mat[h].rank(l, 0), R = mat[h].rank(r, 0);
                res += R - L;
            }
            l = nxt(l, h, a);
            r = nxt(r, h, a);
        }
        return res;
    }
    // [l,r) で [lower,upper) の個数
    int range_freq(int l, int r, const T &lower, const T &upper) {
        return range_freq(l, r, upper) - range_freq(l, r, lower);
    }

    std::optional<T> lt(int l, int r, const T &x) {
        int cnt = range_freq(l, r, x);
        if (cnt)
            return kth_smallest(l, r, cnt - 1);
        return std::nullopt;
    }
    std::optional<T> leq(int l, int r, const T &x) {
        if (rank(l, r, x))
            return x;
        return lt(l, r, x);
    }
    std::optional<T> geq(int l, int r, const T &x) {
        int cnt = r - l - range_freq(l, r, x);
        if (cnt)
            return kth_largest(l, r, cnt - 1);
        return std::nullopt;
    }
    std::optional<T> gt(int l, int r, const T &x) {
        T y;
        if constexpr (COMPRESS) {
            y = C.r(C.gt(x));
        } else {
            y = x + 1;
        }
        return geq(l, r, y);
    }

    // セグ木などを載せる時用
    // BIT は専用のライブラリを作ってあるが、一応抽象化も持っておく
    // 構築したものを返してるので遅そう
    int height() const { return log; }
    std::vector<int> points(int idx) {
        std::vector<int> res(log);
        U a = comp(data[idx]);
        REP_(h, log) {
            idx = nxt(idx, h, a);
            res[h] = idx;
        }
        return res;
    }
    std::vector<std::tuple<int, int, int>> intervals(int l, int r,
                                                     const T &upper) {
        std::vector<std::tuple<int, int, int>> res;
        U a = comp(upper);
        REP_(h, log) {
            if (high_bit(a, h)) {
                int L = mat[h].rank(l, 0), R = mat[h].rank(r, 0);
                res.emplace_back(h, L, R);
            }
            l = nxt(l, h, a);
            r = nxt(r, h, a);
        }
        return res;
    }
    std::vector<std::tuple<int, int, int>> kth_largest_intervals(int l, int r,
                                                                 int k) {
        assert(0 <= k and k < r - l);
        std::vector<std::tuple<int, int, int>> res;
        REP_(h, log) {
            int L = mat[h].rank(l);
            int R = mat[h].rank(r);
            if (R - L > k) {
                l = L + zero_cnt[h];
                r = R + zero_cnt[h];
            } else {
                res.emplace_back(h, L + zero_cnt[h], R + zero_cnt[h]);
                k -= R - L;
                l -= L;
                r -= R;
            }
        }
        return res;
    }
    std::vector<std::tuple<int, int, int>> kth_smallest_intervals(int l, int r,
                                                                  int k) {
        return kth_largest_intervals(l, r, r - l - k - 1);
    }
};
#undef REP_
#line 2 "library/r2/XY.hpp"
template <typename T> struct XY {
    T x, y;
    XY() = default;
    XY(T x, T y) : x(x), y(y) {}
    XY(const std::pair<T, T> &xy) : x(xy.first), y(xy.second) {}

    XY operator+() const { return *this; }
    XY operator-() const { return XY(-x, -y); }

    XY &operator++() {
        x++;
        y++;
        return *this;
    }
    XY &operator--() {
        x--;
        y--;
        return *this;
    }
    XY &operator++(int) {
        XY a = *this;
        ++*this;
        return a;
    }
    XY &operator--(int) {
        XY a = *this;
        --*this;
        return a;
    }

    XY &operator+=(const XY &v) {
        x += v.x;
        y += v.y;
        return *this;
    }
    XY &operator-=(const XY &v) {
        x -= v.x;
        y -= v.y;
        return *this;
    }
    XY &operator*=(const T &a) {
        x *= a;
        y *= a;
        return *this;
    }
    XY &operator/=(const T &a) {
        x /= a;
        y /= a;
        return *this;
    }

    friend XY operator+(const XY &u, const XY &v) { return XY(u) += v; }
    friend XY operator-(const XY &u, const XY &v) { return XY(u) -= v; }
    friend XY operator*(const XY &u, const T &a) { return XY(u) *= a; }
    friend XY operator*(const T &a, const XY &u) { return XY(u) *= a; }
    friend XY operator/(const XY &u, const T &a) { return XY(u) /= a; }

    bool operator<(const XY &v) const { return x != v.x ? x < v.x : y < v.y; }
    bool operator>(const XY &v) const { return x != v.x ? x > v.x : y > v.y; }
    bool operator==(const XY &v) const { return x == v.x and y == v.y; }
    bool operator<=(const XY &v) const { return !(*this > v); }
    bool operator>=(const XY &v) const { return !(*this < v); }
    bool operator!=(const XY &v) const { return !(*this == v); }

    double arg() const { return atan2(y, x); }

    // [0,2pi) で θ(u)<θ(v) の時 true
    // (0,0) は 2pi に相当
    // static bool angle_cmp(const XY&u,const XY&v){
    //  using U=conditional_t< is_same_v<T,int>,long long,T>;
    //  if(u==XY(0,0))return false;
    //  if(v==XY(0,0))return true;
    //  if(u.y==0){
    //    if(u.x>0)return true;
    //    if(v.y==0)return v.x<0;
    //    return v.y<0;
    //  }
    //  if(u.y>0){
    //    if(v.y==0)return v.x<0;
    //    if(v.y<0)return true;
    //    return U(v.x)*u.y <= U(u.x)*v.y;
    //  }
    //  if(v.y>=0)return false;
    //  return U(v.x)*u.y <= U(u.x)*v.y;
    //}

    friend T dot(const XY &u, const XY &v) { return u.x * v.x + u.y * v.y; }
    T norm() { return dot(*this, *this); }
    T abs() { return sqrt(norm()); }

    friend std::istream &operator>>(std::istream &is, XY &v) {
        is >> v.x >> v.y;
        return is;
    }
    friend std::ostream &operator<<(std::ostream &os, const XY &v) {
        os << v.x << " " << v.y;
        return os;
    }

    static XY direction(const char &c) {
        if (c == 'R')
            return {1, 0};
        if (c == 'L')
            return {-1, 0};
        if (c == 'U')
            return {0, -1};
        if (c == 'D')
            return {0, 1};
        return {0, 0};
    }
};
#line 4 "library/r2/Projection.hpp"
template <typename T> class Projection {
    using r2 = XY<T>;
    Compress<r2> C;

  public:
    Projection(const std::vector<r2> &v) : C(v) {}
    int size() { return C.size(); }
    int id(r2 xy) { return C[xy]; }
    int id(int x, int y) { return C[r2(x, y)]; }
    r2 r(int id) { return C.r(id); }
    //[l,r) を返す
    std::pair<int, int> interval(const T &l, const T &r) {
        if (C.max().x < l or r <= C.min().x)
            return std::make_pair(0, 0);
        T mn = std::numeric_limits<T>::min();
        int L = C.geq(r2(l, mn));
        int R = C.geq(r2(r, mn));
        return std::make_pair(L, R);
    }
};
#line 10 "test/library-checker/DataStructure/RectangleSum_2.test.cpp"

using ll = long long;
using r2 = XY<ll>;

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

    int n, q;
    std::cin >> n >> q;
    std::vector<r2> r2s(n);
    std::vector<int> x(n), y(n), w(n);
    REP (i, n) {
        std::cin >> x[i] >> y[i] >> w[i];
        r2s[i] = r2(x[i], y[i]);
    }

    auto P = Projection(r2s);
    std::vector<int> Y(P.size(), 0);
    REP (id, P.size())
        Y[id] = P.r(id).y;
    WaveletMatrix WM(Y);

    int LOG = WM.height();
    std::vector W(LOG, std::vector<ll>(P.size(), 0));
    REP (i, n) {
        int idx = P.id(x[i], y[i]);
        auto idxs = WM.points(idx);
        REP (h, LOG)
            W[h][idxs[h]] += w[i];
    }
    std::vector<CumulativeSum<ll>> C(LOG);
    REP (h, LOG)
        C[h] = CumulativeSum<ll>(W[h]);

    REP (j, q) {
        int l, r, d, u;
        std::cin >> l >> d >> r >> u;
        auto [L, R] = P.interval(l, r);
        ll res = 0;
        auto intervals = WM.intervals(L, R, u);
        for (const auto &[h, l, r] : intervals)
            res += C[h].sum(l, r);
        auto intervals2 = WM.intervals(L, R, d);
        for (const auto &[h, l, r] : intervals2)
            res -= C[h].sum(l, r);
        std::cout << res << "\n";
    }
}

Back to top page