tree/test/lca_rmq.test.cpp

• View this file on GitHub
• Last update: 2022-09-24 13:25:19+09:00
• Problem: https://judge.yosupo.jp/problem/lca

Depends on

• sparse_table/rmq_sparse_table.hpp
• Lowest common ancestor of tree based on sparse table （クエリ $O(1)$ の最小共通祖先） (tree/lca_rmq.hpp)

Code

#include "../lca_rmq.hpp"
#include <iostream>
#define PROBLEM "https://judge.yosupo.jp/problem/lca"
using namespace std;

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

    int N, Q, p, u, v;
    cin >> N >> Q;
    TreeLCA tree(N);
    for (int i = 1; i <= N - 1; i++) cin >> p, tree.add_edge(i, p);

    tree.build(0);

    while (Q--) cout << tree.lca((cin >> u, u), (cin >> v, v)) << '\n';
}

#line 2 "sparse_table/rmq_sparse_table.hpp"
#include <algorithm>
#include <cassert>
#include <vector>

// CUT begin
// Range Minimum Query for static sequence by sparse table
// Complexity: $O(N \log N)$ for precalculation, $O(1)$ per query
template <typename T> struct StaticRMQ {
    inline T func(const T &l, const T &r) const noexcept { return std::min<T>(l, r); }
    int N, lgN;
    T defaultT;
    std::vector<std::vector<T>> data;
    std::vector<int> lgx_table;
    StaticRMQ() = default;
    StaticRMQ(const std::vector<T> &sequence, T defaultT)
        : N(sequence.size()), defaultT(defaultT) {
        lgx_table.resize(N + 1);
        for (int i = 2; i < N + 1; i++) lgx_table[i] = lgx_table[i >> 1] + 1;
        lgN = lgx_table[N] + 1;
        data.assign(lgN, std::vector<T>(N, defaultT));
        data[0] = sequence;
        for (int d = 1; d < lgN; d++) {
            for (int i = 0; i + (1 << d) <= N; i++) {
                data[d][i] = func(data[d - 1][i], data[d - 1][i + (1 << (d - 1))]);
            }
        }
    }
    T get(int l, int r) const { // [l, r), 0-indexed
        assert(l >= 0 and r <= N);
        if (l >= r) return defaultT;
        int d = lgx_table[r - l];
        return func(data[d][l], data[d][r - (1 << d)]);
    }
};
#line 5 "tree/lca_rmq.hpp"
#include <utility>
#line 7 "tree/lca_rmq.hpp"

struct TreeLCA {
    const int N;
    std::vector<std::vector<int>> to;
    int root;
    TreeLCA(int V = 0) : N(V), to(V), root(-1) {}

    void add_edge(int u, int v) {
        assert(0 <= u and u < N);
        assert(0 <= v and v < N);
        assert(u != v);
        to[u].push_back(v);
        to[v].push_back(u);
    }

    using P = std::pair<int, int>;
    std::vector<int> subtree_begin;
    std::vector<P> vis_order;
    std::vector<int> depth;
    void _build_dfs(int now, int prv, int dnow) {
        subtree_begin[now] = vis_order.size();
        vis_order.emplace_back(dnow, now);
        depth[now] = dnow;
        for (auto &&nxt : to[now]) {
            if (nxt != prv) {
                _build_dfs(nxt, now, dnow + 1);
                vis_order.emplace_back(dnow, now);
            }
        }
    }

    StaticRMQ<P> rmq;
    void build(int root_) {
        assert(root_ >= 0 and root_ < N);
        if (root == root_) return;
        root = root_;
        subtree_begin.assign(N, 0);
        vis_order.clear();
        vis_order.reserve(N * 2);
        depth.assign(N, 0);
        _build_dfs(root, -1, 0);
        rmq = {vis_order, P{N, -1}};
    }

    bool built() const noexcept { return root >= 0; }

    int lca(int u, int v) const {
        assert(0 <= u and u < N);
        assert(0 <= v and v < N);
        assert(built());

        int a = subtree_begin[u], b = subtree_begin[v];
        if (a > b) std::swap(a, b);
        return rmq.get(a, b + 1).second;
    };

    int path_length(int u, int v) const { return depth[u] + depth[v] - depth[lca(u, v)] * 2; }
};
#line 2 "tree/test/lca_rmq.test.cpp"
#include <iostream>
#define PROBLEM "https://judge.yosupo.jp/problem/lca"
using namespace std;

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

    int N, Q, p, u, v;
    cin >> N >> Q;
    TreeLCA tree(N);
    for (int i = 1; i <= N - 1; i++) cin >> p, tree.add_edge(i, p);

    tree.build(0);

    while (Q--) cout << tree.lca((cin >> u, u), (cin >> v, v)) << '\n';
}

Back to top page