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

• View this file on GitHub
• Last update: 2022-09-24 13:25:19+09:00
• Include: #include "tree/lca_rmq.hpp"

根を固定した木に対する 2 頂点の最小共通祖先，および 2 頂点間の距離の計算．前処理 $O(N \log N)$, 空間 $O(N \log N)$, クエリ $O(1)$．

使用方法

add_edge(int u, int v) を $N - 1$ 回行って木を構築した後，build(int root) によって根を指定する．

TreeLCA tree(N);
for (int e = 0; e < N - 1; e++) {
    int u, v;
    cin >> u >> v;
    tree.add_edge(u, v);
}

tree.build(0);

cout << tree.lca(a, b) << '\n'; // (a, b) の最長共通祖先
cout << tree.path_length(a, b) << '\n'; // 2 頂点 a, b の距離

問題例

• Library Checker: Lowest Common Ancestor

Depends on

• sparse_table/rmq_sparse_table.hpp

Verified with

• tree/test/lca_rmq.test.cpp

Code

#pragma once
#include "../sparse_table/rmq_sparse_table.hpp"
#include <algorithm>
#include <cassert>
#include <utility>
#include <vector>

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 "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; }
};

Back to top page