DataStructure/LCA.h

• View this file on GitHub
• Last update: 2022-08-11 17:15:37+08:00
• Link: https://judge.yosupo.jp/problem/lca

Depends on

• DataStructure/RMQ.h

Verified with

• DataStructure/test/aizu_grl_5_c_lca.test.cpp
• DataStructure/test/lca.test.cpp

Code

// LCA
// Notes:
// - Index from 0
// - Cannot use for queries like min edge in path u -> v
//
// Tested:
// - https://judge.yosupo.jp/problem/lca
#include "RMQ.h"
struct LCA {
    const int n;
    vector<vector<int>> adj;
    const int root;
    using P = pair<int,int>;
    static P f(P x, P y) { return std::min(x, y); }

    LCA(int _n, const vector<vector<int>>& _adj, int _root)
            : n(_n), adj(_adj), root(_root) {
        assert(0 <= root && root < n);
        id.resize(n);
        depth.resize(n);
        order.reserve(2 * n);
        dfs(root, -1, 0);
        rmq = RMQ<P, f>(order);
    }

    int lca(int u, int v) {
        assert(0 <= u && u < n);
        assert(0 <= v && v < n);

        int x = id[u], y = id[v];
        if (x > y) std::swap(x, y);
        return rmq.get(x, y+1).second;
    }

// private:
    vector<int> id, depth;
    vector<P> order;
    RMQ<P, f> rmq;

    void dfs(int u, int fu, int d) {
        id[u] = order.size();
        depth[u] = d;
        order.emplace_back(d, u);
        for (int v : adj[u]) {
            if (v == fu) continue;
            dfs(v, u, d + 1);
            order.emplace_back(d, u);
        }
    }
};

#line 1 "DataStructure/LCA.h"
// LCA
// Notes:
// - Index from 0
// - Cannot use for queries like min edge in path u -> v
//
// Tested:
// - https://judge.yosupo.jp/problem/lca
#line 1 "DataStructure/RMQ.h"
// RMQ {{{
//
// Sparse table
// Usage:
// RMQ<int, _min> st(v);
//
// Note:
// - doesn't work for empty range
//
// Tested:
// - https://judge.yosupo.jp/problem/staticrmq
template<class T, T (*op) (T, T)> struct RMQ {
    RMQ() = default;
    RMQ(const vector<T>& v) : t{v}, n{(int) v.size()} {
        for (int k = 1; (1<<k) <= n; ++k) {
            t.emplace_back(n - (1<<k) + 1);
            for (int i = 0; i + (1<<k) <= n; ++i) {
                t[k][i] = op(t[k-1][i], t[k-1][i + (1<<(k-1))]);
            }
        }
    }

    // get range [l, r-1]
    // doesn't work for empty range
    T get(int l, int r) const {
        assert(0 <= l && l < r && r <= n);
        int k = __lg(r - l);
        return op(t[k][l], t[k][r - (1<<k)]);
    }

private:
    vector<vector<T>> t;
    int n;
};
template<class T> T _min(T a, T b) { return b < a ? b : a; }
template<class T> T _max(T a, T b) { return a < b ? b : a; }
// }}}
#line 9 "DataStructure/LCA.h"
struct LCA {
    const int n;
    vector<vector<int>> adj;
    const int root;
    using P = pair<int,int>;
    static P f(P x, P y) { return std::min(x, y); }

    LCA(int _n, const vector<vector<int>>& _adj, int _root)
            : n(_n), adj(_adj), root(_root) {
        assert(0 <= root && root < n);
        id.resize(n);
        depth.resize(n);
        order.reserve(2 * n);
        dfs(root, -1, 0);
        rmq = RMQ<P, f>(order);
    }

    int lca(int u, int v) {
        assert(0 <= u && u < n);
        assert(0 <= v && v < n);

        int x = id[u], y = id[v];
        if (x > y) std::swap(x, y);
        return rmq.get(x, y+1).second;
    }

// private:
    vector<int> id, depth;
    vector<P> order;
    RMQ<P, f> rmq;

    void dfs(int u, int fu, int d) {
        id[u] = order.size();
        depth[u] = d;
        order.emplace_back(d, u);
        for (int v : adj[u]) {
            if (v == fu) continue;
            dfs(v, u, d + 1);
            order.emplace_back(d, u);
        }
    }
};

Back to top page