DataStructure/test/mo_algorithm_with_undo.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/MoAlgorithmWithUndo.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;
    stack<int> nums;

    auto add = [&] (int id) {
        sum += a[id];
        nums.push(a[id]);
    };
    auto undo = [&] () {
        sum -= nums.top();
        nums.pop();
    };
    auto get = [&] ([[maybe_unused]] const Query& _) {
        return sum;
    };

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

#line 1 "DataStructure/test/mo_algorithm_with_undo.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/MoAlgorithmWithUndo.h"
// Notes:
// - queries are [l, r]
// - add(int array_id) -> void 
// - undo() -> void
// - get(QueryT query) -> ResultT
//
// Tested:
// - https://www.spoj.com/problems/FREQ2/ (submission ID: 30602190)
// - https://www.spoj.com/problems/KDOMINO/ (submission ID: 30602247)
// - https://www.spoj.com/problems/ZQUERY/
//
// Mo algorithm with undo DS {{{
template<typename QueryT, typename ResultT, typename Add, typename Undo, typename Get>
vector<ResultT> mo_with_undo(int n, std::vector<QueryT> queries, Add add, Undo undo, Get get) {
    int s = max((int) 1, (int) sqrt(n + 0.5));
    int q = queries.size();
    vector<ResultT> res(q);

    // For each block, stores vector of (query, id)
    std::vector<std::vector<int>> blocks(n / s + 1);
    for (int i = 0; i < q; ++i) {
        blocks[queries[i].l / s].push_back(i);
    }

    // Process each block
    int block_start = 0, block_end = s - 1;
    for (auto& block : blocks) {
        sort(block.begin(), block.end(), [&] (int q1, int q2) {
                return queries[q1].r < queries[q2].r;
                });
        int cur_r = block_end;
        for (int qid : block) {
            const auto& query = queries[qid];
            assert(block_start <= query.l && query.l <= block_end);

            // For query that is entirely inside this block, we just process it
            // and then undo immediately
            if (query.r <= block_end) {
                for (int i = query.l; i <= query.r; ++i) add(i);
                res[qid] = get(query);
                for (int i = query.l; i <= query.r; ++i) undo();
                continue;
            }

            // Order:
            // 1. add(i) for i in block_end + 1 .. query.r
            // 2. add(i) for i in query.l .. block_end
            // 3. get()
            // 4. undo() for i in query.l .. block_end => all of (2) is undo
            for (int i = cur_r + 1; i <= query.r; ++i) add(i);
            for (int i = query.l; i <= block_end; ++i) add(i);
            res[qid] = get(query);
            for (int i = query.l; i <= block_end; ++i) undo();
            cur_r = query.r;
        }

        for (int i = cur_r; i > block_end; --i) undo();

        block_start += s;
        block_end += s;
    }
    return res;
}

// Example
struct Query {
    int l, r;  // QueryT must have l, r
};
// Usage
// auto res = mo_with_undo<Query, int, decltype(add), decltype(undo), decltype(get)>
//        (n, queries, add, undo, get);
// }}}

#line 5 "DataStructure/test/mo_algorithm_with_undo.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;
    stack<int> nums;

    auto add = [&] (int id) {
        sum += a[id];
        nums.push(a[id]);
    };
    auto undo = [&] () {
        sum -= nums.top();
        nums.pop();
    };
    auto get = [&] ([[maybe_unused]] const Query& _) {
        return sum;
    };

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