test/mo/aoj0489.test.cpp

View this file on GitHub
Last update: 2026-02-28 01:07:59+09:00
Problem: https://onlinejudge.u-aizu.ac.jp/problems/0489

Depends on

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/problems/0489"

#include "mo/mo_tree.hpp"
#include "mo/sqrt_freq_table.hpp"
#include <bits/stdc++.h>

int main() {
    std::cin.tie(0)->sync_with_stdio(0);
    int N, M;
    std::cin >> N >> M;
    std::vector<int> A(N);
    for (auto &x : A) {
        std::cin >> x;
    }
    std::vector<std::vector<int>> adj(N);
    for (auto i = 0; i < N - 1; ++i) {
        int u, v;
        std::cin >> u >> v;
        --u, --v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    std::vector<std::tuple<int, int, int>> queries;
    for (auto i = 0; i < M; ++i) {
        int op;
        std::cin >> op;
        if (op == 1) {
            int u, w;
            std::cin >> u >> w;
            --u;
            auto v = int(adj.size());
            A.push_back(w);
            adj.emplace_back();
            adj[u].push_back(v);
            adj[v].push_back(u);
        } else {
            int u, v, k;
            std::cin >> u >> v >> k;
            --u, --v, --k;
            queries.emplace_back(u, v, k);
        }
    }
    auto C = A;
    std::ranges::sort(C);
    C.erase(unique(C.begin(), C.end()), C.end());
    for (auto &x : A) {
        auto it = std::ranges::lower_bound(C, x);
        x = int(it - C.begin());
    }
    MoTree mo(adj);
    for (auto [u, v, k] : queries) {
        mo.add(u, v);
    }
    SqrtFreqTable table(int(C.size()) - 1);
    std::vector<int> ans(queries.size());
    auto add = [&](int u) { table.insert(A[u]); };
    auto remove = [&](int u) { table.erase(A[u]); };
    auto eval = [&](int i) {
        auto [u, v, k] = queries[i];
        ans[i] = C[table.kth(k)];
    };
    mo.solve(add, remove, eval);
    std::ranges::copy(ans, std::ostream_iterator<int>(std::cout, "\n"));
}

#line 1 "test/mo/aoj0489.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/problems/0489"

#line 1 "mo/mo_tree.hpp"



#line 1 "mo/hilbert_mo.hpp"



#include <algorithm>
#include <array>
#include <cassert>
#include <vector>

struct HilbertMo {
    HilbertMo(int n) : n_(n), log_(std::bit_width(unsigned(n))) {}
    void add(int l, int r) {
        assert(0 <= l && l <= r && r <= n_);
        auto index = int(queries_.size());
        auto order = hilbert_order(l, r);
        queries_.push_back({l, r, index, order});
    }
    void solve(auto add, auto remove, auto eval) { solve(add, add, remove, remove, eval); }
    void solve(auto add_left, auto add_right, auto remove_left, auto remove_right, auto eval) {
        sort(queries_.begin(), queries_.end());
        auto l = 0, r = 0;
        for (auto [left, right, index, order] : queries_) {
            while (left < l) {
                add_left(--l);
            }
            while (r < right) {
                add_right(r++);
            }
            while (l < left) {
                remove_left(l++);
            }
            while (right < r) {
                remove_right(--r);
            }
            eval(index);
        }
    }

private:
    struct query {
        int left, right, index;
        long long order;
        bool operator<(const query &other) const { return order < other.order; }
    };
    long long hilbert_order(int x, int y) const {
        auto d = 0LL;
        for (auto s = 1 << log_; s > 0; s >>= 1) {
            bool rx = x & s, ry = y & s;
            d = (d << 2) | ((rx * 3) ^ ry);
            if (!ry) {
                if (rx) {
                    x = ~x;
                    y = ~y;
                }
                std::swap(x, y);
            }
        }
        return d;
    }
    std::vector<query> queries_;
    const int n_, log_;
};


#line 7 "mo/mo_tree.hpp"
#include <utility>
#line 9 "mo/mo_tree.hpp"

struct MoTree {
    explicit MoTree(const std::vector<std::vector<int>> &adj)
        : n_(int(adj.size())), in_(n_), tour_(2 * n_), top_(n_), parent_(n_), mo_(2 * n_) {
        auto timer = 0;
        std::vector<int> size(n_, 1), heavy(n_, -1);
        auto dfs = [&](auto &&self, int u) -> void {
            for (auto v : adj[u]) {
                if (v != parent_[u]) {
                    parent_[v] = u;
                    self(self, v);
                    size[u] += size[v];
                    if (heavy[u] == -1 || size[heavy[u]] < size[v]) {
                        heavy[u] = v;
                    }
                }
            }
        };
        auto hld = [&](auto &&self, int u) -> void {
            in_[u] = timer;
            tour_[timer++] = u;
            for (auto v : adj[u]) {
                if (v != parent_[u]) {
                    top_[v] = (v == heavy[u]) ? top_[u] : v;
                    self(self, v);
                }
            }
            tour_[timer++] = u;
        };
        dfs(dfs, 0);
        hld(hld, 0);
    }
    int lca(int u, int v) const {
        assert(0 <= u && u < n_ && 0 <= v && v <= n_);
        while (top_[u] != top_[v]) {
            if (in_[top_[u]] < in_[top_[v]]) {
                std::swap(u, v);
            }
            u = parent_[top_[u]];
        }
        return (in_[u] < in_[v]) ? u : v;
    }
    void add(int u, int v) {
        assert(0 <= u && u < n_ && 0 <= v && v <= n_);
        lca_.push_back(lca(u, v));
        auto [l, r] = std::minmax(in_[u], in_[v]);
        mo_.add(l + 1, r + 1);
    }
    void solve(auto add, auto remove, auto eval) {
        std::vector<bool> contains(n_);
        auto toggle = [&](int i) {
            auto u = tour_[i];
            if (contains[u]) {
                remove(u);
            } else {
                add(u);
            }
            contains[u].flip();
        };
        auto eval_lca = [&](int i) {
            toggle(in_[lca_[i]]);
            eval(i);
            toggle(in_[lca_[i]]);
        };
        mo_.solve(toggle, toggle, eval_lca);
    }

private:
    int n_;
    std::vector<int> in_, tour_, top_, parent_, lca_;
    HilbertMo mo_;
};


#line 1 "mo/sqrt_freq_table.hpp"



#line 6 "mo/sqrt_freq_table.hpp"
#include <cmath>
#line 8 "mo/sqrt_freq_table.hpp"

struct SqrtFreqTable {
    explicit SqrtFreqTable(int max_val)
        : m_(std::max(1, max_val)), b_sz_(int(std::sqrt(m_))), freq_(m_ + 1),
          bucket_(m_ / b_sz_ + 1) {}
    void insert(int x) { // O(1)
        assert(0 <= x && x <= m_);
        ++freq_[x];
        ++bucket_[x / b_sz_];
        ++total_;
        distinct_ += int(freq_[x] == 1);
    }
    void erase(int x) { // O(1)
        assert(0 <= x && x <= m_);
        assert(0 < freq_[x]);
        --freq_[x];
        --bucket_[x / b_sz_];
        --total_;
        distinct_ -= int(freq_[x] == 0);
    }
    int count(int x) const { // O(1)
        assert(0 <= x && x <= m_);
        return freq_[x];
    }
    int size() const { return total_; }
    int count_distinct() const { return distinct_; }
    int kth(int k) const { // O(sqrt(M)), 0-indexed
        auto cnt = 0;
        for (auto i = 0; i < m_ / b_sz_ + 1; ++i) {
            if (cnt + bucket_[i] <= k) {
                cnt += bucket_[i];
                continue;
            }
            auto x = i * b_sz_;
            while ((cnt += freq_[x]) <= k) {
                ++x;
            }
            return x;
        }
        return -1;
    };
    int rank(int x) const { // O(sqrt(M)), count y s.t. y < x
        if (x < 0) {
            return 0;
        }
        if (m_ < x) {
            return total_;
        }
        auto cnt = 0;
        for (auto i = 0; i < x / b_sz_; ++i) {
            cnt += bucket_[i];
        }
        for (auto i = x / b_sz_ * b_sz_; i < x; ++i) {
            cnt += freq_[i];
        }
        return cnt;
    }

private:
    const int m_, b_sz_;
    std::vector<int> freq_, bucket_;
    int total_ = 0, distinct_ = 0;
};


#line 5 "test/mo/aoj0489.test.cpp"
#include <bits/stdc++.h>

int main() {
    std::cin.tie(0)->sync_with_stdio(0);
    int N, M;
    std::cin >> N >> M;
    std::vector<int> A(N);
    for (auto &x : A) {
        std::cin >> x;
    }
    std::vector<std::vector<int>> adj(N);
    for (auto i = 0; i < N - 1; ++i) {
        int u, v;
        std::cin >> u >> v;
        --u, --v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    std::vector<std::tuple<int, int, int>> queries;
    for (auto i = 0; i < M; ++i) {
        int op;
        std::cin >> op;
        if (op == 1) {
            int u, w;
            std::cin >> u >> w;
            --u;
            auto v = int(adj.size());
            A.push_back(w);
            adj.emplace_back();
            adj[u].push_back(v);
            adj[v].push_back(u);
        } else {
            int u, v, k;
            std::cin >> u >> v >> k;
            --u, --v, --k;
            queries.emplace_back(u, v, k);
        }
    }
    auto C = A;
    std::ranges::sort(C);
    C.erase(unique(C.begin(), C.end()), C.end());
    for (auto &x : A) {
        auto it = std::ranges::lower_bound(C, x);
        x = int(it - C.begin());
    }
    MoTree mo(adj);
    for (auto [u, v, k] : queries) {
        mo.add(u, v);
    }
    SqrtFreqTable table(int(C.size()) - 1);
    std::vector<int> ans(queries.size());
    auto add = [&](int u) { table.insert(A[u]); };
    auto remove = [&](int u) { table.erase(A[u]); };
    auto eval = [&](int i) {
        auto [u, v, k] = queries[i];
        ans[i] = C[table.kth(k)];
    };
    mo.solve(add, remove, eval);
    std::ranges::copy(ans, std::ostream_iterator<int>(std::cout, "\n"));
}