DataStructure/test/mo_algorithm.test.cpp

View this file on GitHub
Last update: 2023-10-03 00:47:04-07:00
Problem: https://judge.yosupo.jp/problem/static_range_sum

Depends on

Code

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

#include "../../template.h"
#include "../Mo/MoAlgorithm.h"

void solve() {
    int n, q; cin >> n >> q;
    vector<int> a(n); REP(i,n) cin >> a[i];

    vector<Query> queries(q);
    REP(i,q) {
        cin >> queries[i].l >> queries[i].r;
        queries[i].r--;
    }

    int64_t sum = 0;
    auto add = [&] (int id) { sum += a[id]; };
    auto rem = [&] (int id) { sum -= a[id]; };
    auto get = [&] ([[maybe_unused]] const Query& _) { return sum; };

    auto res = mo<Query, int64_t, decltype(add), decltype(rem), decltype(get)>
        (n, queries, add, rem, get);
    for (auto r : res) cout << r << '\n';
}

#line 1 "DataStructure/test/mo_algorithm.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/static_range_sum"

#line 1 "template.h"
#include <bits/stdc++.h>
using namespace std;

#define FOR(i,a,b) for(int i=(a),_b=(b); i<=_b; i++)
#define FORD(i,a,b) for(int i=(a),_b=(b); i>=_b; i--)
#define REP(i,a) for(int i=0,_a=(a); i<_a; i++)
#define EACH(it,a) for(__typeof(a.begin()) it = a.begin(); it != a.end(); ++it)

#define DEBUG(x) { cout << #x << " = "; cout << (x) << endl; }
#define PR(a,n) { cout << #a << " = "; FOR(_,1,n) cout << a[_] << ' '; cout << endl; }
#define PR0(a,n) { cout << #a << " = "; REP(_,n) cout << a[_] << ' '; cout << endl; }

#define sqr(x) ((x) * (x))

// For printing pair, container, etc.
// Copied from https://quangloc99.github.io/2021/07/30/my-CP-debugging-template.html
template<class U, class V> ostream& operator << (ostream& out, const pair<U, V>& p) {
    return out << '(' << p.first << ", " << p.second << ')';
}

template<class Con, class = decltype(begin(declval<Con>()))>
typename enable_if<!is_same<Con, string>::value, ostream&>::type
operator << (ostream& out, const Con& con) {
    out << '{';
    for (auto beg = con.begin(), it = beg; it != con.end(); it++) {
        out << (it == beg ? "" : ", ") << *it;
    }
    return out << '}';
}
template<size_t i, class T> ostream& print_tuple_utils(ostream& out, const T& tup) {
    if constexpr(i == tuple_size<T>::value) return out << ")"; 
    else return print_tuple_utils<i + 1, T>(out << (i ? ", " : "(") << get<i>(tup), tup); 
}
template<class ...U> ostream& operator << (ostream& out, const tuple<U...>& t) {
    return print_tuple_utils<0, tuple<U...>>(out, t);
}

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
long long get_rand(long long r) {
    return uniform_int_distribution<long long> (0, r-1)(rng);
}

template<typename T>
vector<T> read_vector(int n) {
    vector<T> res(n);
    for (int& x : res) cin >> x;
    return res;
}

void solve();

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    solve();
    return 0;
}
#line 1 "DataStructure/Mo/MoAlgorithm.h"
// Notes:
// - queries are [l, r]
// - add(int array_id) -> void 
// - remove(int array_id) -> void
// - get(QueryT query) -> ResultT
//
// Tested:
// - https://www.spoj.com/problems/KDOMINO/ (submission ID: 30602374)
// - https://www.spoj.com/problems/FREQ2 (submission ID: 30602401)
//
// Mo algorithm {{{
template<typename QueryT, typename ResultT, typename Add, typename Rem, typename Get>
vector<ResultT> mo(int n, std::vector<QueryT> queries, Add add, Rem rem, Get get) {
    int q = queries.size();
    std::vector<ResultT> res(q);
 
    // Sort queries in increasing order of (left / SQRT, right)
    int S = sqrt(n);
    if (S < 1) S = 1;
 
    std::vector<int> query_ids(q);
    std::iota(query_ids.begin(), query_ids.end(), 0);
    std::sort(query_ids.begin(), query_ids.end(), [&] (int q1, int q2) {
            int bucket1 = queries[q1].l / S;
            int bucket2 = queries[q2].l / S;
            if (bucket1 != bucket2) return bucket1 < bucket2;
            else {
                return bucket1 % 2
                        ? (queries[q1].r > queries[q2].r)
                        : (queries[q1].r < queries[q2].r);
            }
        });
 
    // Answer queries
    int cur_l = -1, cur_r = -1;
    for (int qid : query_ids) {
        // move to this range
        if (cur_l < 0) {
            for (int i = queries[qid].l; i <= queries[qid].r; ++i) {
                add(i);
            }
            cur_l = queries[qid].l, cur_r = queries[qid].r;
        } else {
            while (cur_l > queries[qid].l) add(--cur_l);
            while (cur_r < queries[qid].r) add(++cur_r);
            while (cur_r > queries[qid].r) rem(cur_r--);
            while (cur_l < queries[qid].l) rem(cur_l++);
        }
 
        res[qid] = get(queries[qid]);
    }
    return res;
}
 
// Example
struct Query {
    int l, r;  // QueryT must have l, r
};
ostream& operator << (ostream& cout, const Query& q) {
    cout << "Query: [" << q.l << ", " << q.r << "]";
    return cout;
}
// Usage
// auto res = mo<Query, int, decltype(add), decltype(rem), decltype(get)>
//        (n, queries, add, rem, get);
// }}}
#line 5 "DataStructure/test/mo_algorithm.test.cpp"

void solve() {
    int n, q; cin >> n >> q;
    vector<int> a(n); REP(i,n) cin >> a[i];

    vector<Query> queries(q);
    REP(i,q) {
        cin >> queries[i].l >> queries[i].r;
        queries[i].r--;
    }

    int64_t sum = 0;
    auto add = [&] (int id) { sum += a[id]; };
    auto rem = [&] (int id) { sum -= a[id]; };
    auto get = [&] ([[maybe_unused]] const Query& _) { return sum; };

    auto res = mo<Query, int64_t, decltype(add), decltype(rem), decltype(get)>
        (n, queries, add, rem, get);
    for (auto r : res) cout << r << '\n';
}