test/unit/mo.test.cpp

• View this file on GitHub
• Last update: 2026-10-26 21:23:55+09:00
• Problem: http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A

Depends on

• Mo's Algorithm (algorithm/mo.hpp)

Code

#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <algorithm>
#include <cassert>
#include <iostream>
#include <random>
#include <vector>

#include "algorithm/mo.hpp"

int main() {
    std::mt19937 rng(12345);
    for (int it = 0; it < 300; it++) {
        int n = rng() % 30, q = rng() % 40;
        std::vector<int> a(n);
        for (auto& x : a) x = rng() % 5;
        mo m;
        std::vector<std::pair<int, int>> qs(q);
        for (int i = 0; i < q; i++) {
            int l = rng() % (n + 1), r = rng() % (n + 1);
            if (l > r) std::swap(l, r);
            qs[i] = {l, r};
            assert(m.add(l, r) == i);
        }
        // distinct count: symmetric callbacks
        {
            std::vector<int> cnt(5), ans(q, -1);
            int distinct = 0;
            m.run([&](int i) { if (cnt[a[i]]++ == 0) distinct++; },
                  [&](int i) { if (--cnt[a[i]] == 0) distinct--; },
                  [&](int k) { ans[k] = distinct; });
            for (int k = 0; k < q; k++) {
                auto [l, r] = qs[k];
                std::vector<int> b(a.begin() + l, a.begin() + r);
                std::sort(b.begin(), b.end());
                assert(ans[k] == (int)(std::unique(b.begin(), b.end()) - b.begin()));
            }
            for (int c : cnt) assert(c >= 0);
        }
        // inversions: left/right differ, and the window must satisfy l <= r
        {
            std::vector<int> cnt(5), ans(q, -1);
            int size = 0, inv = 0;
            auto less = [&](int v) { int s = 0; for (int u = 0; u < v; u++) s += cnt[u]; return s; };
            auto greater = [&](int v) { int s = 0; for (int u = v + 1; u < 5; u++) s += cnt[u]; return s; };
            m.run([&](int i) { inv += less(a[i]); cnt[a[i]]++; size++; },
                  [&](int i) { inv += greater(a[i]); cnt[a[i]]++; size++; },
                  [&](int i) { cnt[a[i]]--; size--; inv -= less(a[i]); assert(size >= 0); },
                  [&](int i) { cnt[a[i]]--; size--; inv -= greater(a[i]); assert(size >= 0); },
                  [&](int k) { ans[k] = inv; });
            for (int k = 0; k < q; k++) {
                auto [l, r] = qs[k];
                int e = 0;
                for (int i = l; i < r; i++)
                    for (int j = i + 1; j < r; j++) e += a[i] > a[j];
                assert(ans[k] == e);
            }
        }
    }
    // no queries
    {
        mo m;
        int calls = 0;
        auto f = [&](int) { calls++; };
        m.run(f, f, f);
        assert(calls == 0);
    }
    std::cout << "Hello World" << std::endl;
}

#line 1 "test/unit/mo.test.cpp"
#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <algorithm>
#include <cassert>
#include <iostream>
#include <random>
#include <vector>

#line 4 "algorithm/mo.hpp"
#include <cmath>
#include <cstdlib>
#include <numeric>
#include <utility>
#line 9 "algorithm/mo.hpp"

struct mo {
    mo() = default;
    int add(int l, int r) {
        assert(0 <= l && l <= r);
        lr.emplace_back(l, r);
        return (int)lr.size() - 1;
    }
    void run(auto add_left, auto add_right, auto erase_left, auto erase_right,
             auto query) const {
        int l = 0, r = 0;
        for (int i : order()) {
            while (l > lr[i].first) add_left(--l);
            while (r < lr[i].second) add_right(r++);
            while (l < lr[i].first) erase_left(l++);
            while (r > lr[i].second) erase_right(--r);
            query(i);
        }
    }
    void run(auto add, auto erase, auto query) const {
        run(add, add, erase, erase, query);
    }

  private:
    std::vector<std::pair<int, int>> lr;
    std::vector<int> order() const {
        int q = lr.size();
        if (q == 0) return {};
        int n = 1;
        for (auto [l, r] : lr) n = std::max(n, r);
        int w = std::max(1, (int)(std::sqrt(3.0) * n / std::sqrt(2.0 * q)));
        std::vector<int> ord(q);
        std::iota(ord.begin(), ord.end(), 0);
        std::sort(ord.begin(), ord.end(), [&](int a, int b) {
            int ba = lr[a].first / w, bb = lr[b].first / w;
            if (ba != bb) return ba < bb;
            return (ba & 1) ? lr[a].second > lr[b].second
                            : lr[a].second < lr[b].second;
        });
        auto cost = [&](int a, int b) {
            return std::abs(lr[ord[a]].first - lr[ord[b]].first) +
                   std::abs(lr[ord[a]].second - lr[ord[b]].second);
        };
        auto improve = [&](int k) {
            if (cost(k, k + 2) + cost(k + 1, k + 3) <
                cost(k, k + 1) + cost(k + 2, k + 3))
                std::swap(ord[k + 1], ord[k + 2]);
            if (cost(k, k + 3) + cost(k + 1, k + 4) <
                cost(k, k + 1) + cost(k + 3, k + 4))
                std::swap(ord[k + 1], ord[k + 3]);
        };
        for (int k = 0; k + 5 <= q; k++) improve(k);
        for (int k = q - 5; k >= 0; k--) improve(k);
        return ord;
    }
};
#line 9 "test/unit/mo.test.cpp"

int main() {
    std::mt19937 rng(12345);
    for (int it = 0; it < 300; it++) {
        int n = rng() % 30, q = rng() % 40;
        std::vector<int> a(n);
        for (auto& x : a) x = rng() % 5;
        mo m;
        std::vector<std::pair<int, int>> qs(q);
        for (int i = 0; i < q; i++) {
            int l = rng() % (n + 1), r = rng() % (n + 1);
            if (l > r) std::swap(l, r);
            qs[i] = {l, r};
            assert(m.add(l, r) == i);
        }
        // distinct count: symmetric callbacks
        {
            std::vector<int> cnt(5), ans(q, -1);
            int distinct = 0;
            m.run([&](int i) { if (cnt[a[i]]++ == 0) distinct++; },
                  [&](int i) { if (--cnt[a[i]] == 0) distinct--; },
                  [&](int k) { ans[k] = distinct; });
            for (int k = 0; k < q; k++) {
                auto [l, r] = qs[k];
                std::vector<int> b(a.begin() + l, a.begin() + r);
                std::sort(b.begin(), b.end());
                assert(ans[k] == (int)(std::unique(b.begin(), b.end()) - b.begin()));
            }
            for (int c : cnt) assert(c >= 0);
        }
        // inversions: left/right differ, and the window must satisfy l <= r
        {
            std::vector<int> cnt(5), ans(q, -1);
            int size = 0, inv = 0;
            auto less = [&](int v) { int s = 0; for (int u = 0; u < v; u++) s += cnt[u]; return s; };
            auto greater = [&](int v) { int s = 0; for (int u = v + 1; u < 5; u++) s += cnt[u]; return s; };
            m.run([&](int i) { inv += less(a[i]); cnt[a[i]]++; size++; },
                  [&](int i) { inv += greater(a[i]); cnt[a[i]]++; size++; },
                  [&](int i) { cnt[a[i]]--; size--; inv -= less(a[i]); assert(size >= 0); },
                  [&](int i) { cnt[a[i]]--; size--; inv -= greater(a[i]); assert(size >= 0); },
                  [&](int k) { ans[k] = inv; });
            for (int k = 0; k < q; k++) {
                auto [l, r] = qs[k];
                int e = 0;
                for (int i = l; i < r; i++)
                    for (int j = i + 1; j < r; j++) e += a[i] > a[j];
                assert(ans[k] == e);
            }
        }
    }
    // no queries
    {
        mo m;
        int calls = 0;
        auto f = [&](int) { calls++; };
        m.run(f, f, f);
        assert(calls == 0);
    }
    std::cout << "Hello World" << std::endl;
}

Back to top page