test/lca.test.cpp

View this file on GitHub
Last update: 2024-07-08 16:24:06+09:00
Problem: https://judge.yosupo.jp/problem/lca

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/lca"
#include "C++/graph/Graph.hpp"
#include "C++/graph/LCA.hpp"
int main() {
    std::cin.tie(nullptr) -> sync_with_stdio(false);
    int n, q;
    std::cin >> n >> q;
    graph<false> g(n, 0);
    for(int i = 1; i < n; ++i) {
        int p;
        std::cin >> p;
        g.add(p, i);
    }
    LowestCommonAncestor lca(g);
    while(q--) {
        int u, v;
        std::cin >> u >> v;
        std::cout << lca.query(u, v) << '\n';
    }
}

#line 1 "test/lca.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/lca"
#line 2 "C++/graph/Graph.hpp"

#include <iostream>

#include <vector>

#include <algorithm>

#include <queue>

#include <stack>

#ifndef TEMPLATE
template <class T, class U> bool chmin(T& a, const U& b){ if(a>b){ a=b; return 1; } return 0; }
#endif
#line 2 "C++/graph/edge.hpp"

struct edge {
    int src, to, id;
    long long cost;
    edge(){}
    edge(const int src_, const int to_, const int id_ = -1, const long long cost_ = 0): src(src_), to(to_), id(id_), cost(cost_){}
    operator int() const { return to; }
};

/**
 * @brief Edge
 */
#line 12 "C++/graph/Graph.hpp"
template <bool undirected = true> struct graph: std::vector<std::vector<edge>> {
protected:
    int indexed, id;
    std::vector<edge> edges;
public:
    graph(){}
    graph(const int n, const int indexed_ = 1): indexed(indexed_), id(0){ this -> resize(n); }
    void add(int a, int b) {
        a -= indexed, b-= indexed;
        (*this)[a].emplace_back(a, b, id);
        edges.emplace_back(a, b, id++);
        if(undirected) {
            (*this)[b].emplace_back(b, a, --id);
            edges.emplace_back(b, a, id++);
        }
    }
    void input(const int m) {
        for(int i = 0; i < m; ++i) {
            int a, b;
            std::cin >> a >> b;
            add(a, b);
        }
    }
    std::vector<edge> get_edge() const { return edges; }
    std::vector<int> all_dist(const int v) {
        std::vector<int> d(this -> size(), -1);
        std::queue<int> q;
        d[v] = 0;
        q.emplace(v);
        while(q.size()) {
            const int tmp = q.front();
            q.pop();
            for(const auto &el: (*this)[tmp]) {
                if(d[el] != -1) {
                    continue;
                }
                d[el] = d[tmp] + 1;
                q.emplace(el);
            }
        }
        return d;
    }
    int dist(const int u, const int v){ return all_dist(u)[v]; }
    std::vector<int> t_sort() {
        const int n = this -> size();
		std::vector<int> deg(n);
		for(int i = 0; i < n; ++i) {
			for(const auto ed: (*this)[i]) {
				deg[ed]++;
			}
		}
		std::stack<int> sk;
		for(int i = 0; i < n; ++i) {
			if(deg[i] == 0) {
				sk.emplace(i);
			}
		}
		std::vector<int> ord;
		while(sk.size()) {
			const auto tmp = sk.top();
            sk.pop();
			ord.emplace_back(tmp);
			for(const auto ed: (*this)[tmp]) {
				if(--deg[ed] == 0) {
					sk.emplace(ed);
				}
			}
		}
		return ord.size() == size() ? ord : std::vector<int>{};
	}
    std::vector<edge> cycle() {
        const int n = size();
        std::vector<int> used(n);
        std::vector<edge> pre(n), cycle;
        const auto dfs = [&](const auto &f, const int i) -> bool {
            used[i] = 1;
			for(const auto &e: (*this)[i]) {
				if(used[e] == 0) {
					pre[e] = e;
					if(f(f, e)) {
						return true;
					}
				} else if(used[e] == 1) {
					int now = i;
					while(now != e) {
						cycle.emplace_back(pre[now]);
						now = pre[now].src;
					}
					cycle.emplace_back(e);
					return true;
				}
			}
			used[i] = 2;
			return false;
        };
        for(int i = 0; i < n; ++i) {
			if(used[i] == 0 && dfs(dfs, i)) {
				std::reverse(cycle.begin(), cycle.end());
				return cycle;
			}
		}
		return {};
    }
};
typedef std::vector<edge> ve;
typedef std::vector<ve> we;

/**
 * @brief グラフライブラリ
 */
#line 2 "C++/graph/LCA.hpp"

#pragma GCC diagnostic ignored "-Wreorder"

#line 6 "C++/graph/LCA.hpp"
template <class G> struct LowestCommonAncestor {
private:
    const int LOG;
    std::vector<int> dep, sum;
    const G &g;
    std::vector<std::vector<int>> table;
    void dfs(const int idx, const int par, const int d) {
        table[0][idx] = par;
        dep[idx] = d;
        for(const auto &el: g[idx]) {
            if(el != par) {
                sum[el] = sum[idx] + el.cost;
                dfs(el, idx, d + 1);
            }
        }
    }
    void build() {
        dfs(0, -1, 0);
        for(int k = 0; k < LOG - 1; ++k) {
            for(size_t i = 0; i < table[k].size(); ++i) {
                if(table[k][i] == -1) {
                    table[k + 1][i] = -1;
                }
                else {
                    table[k + 1][i] = table[k][table[k][i]];
                }
            }
        }
    }
public:
    LowestCommonAncestor(const G &g_): g(g_), dep(g_.size()), sum(g_.size()), LOG(std::__lg(g_.size()) + 1) {
        table.assign(LOG, std::vector<int>(g_.size(), -1));
        build();
    }
    int climb(int u, const int k) {
		if(dep[u] < k) {
			return -1;
		}
		for(int i = LOG; --i >= 0;) {
			if((k >> i) & 1) {
                u = table[i][u];
            }
		}
		return u;
	}
    int query(int u, int v) {
        if(dep[u] > dep[v]) {
            std::swap(u, v);
        }
        v = climb(v, dep[v] - dep[u]);
        if(u == v) {
            return u;
        }
        for(int i = LOG; --i >= 0;) {
            if(table[i][u] != table[i][v]) {
                u = table[i][u];
                v = table[i][v];
            }
        }
        return table[0][u];
    }
    int dist(const int u, const int v){ return sum[u] + sum[v] - 2 * sum[query(u, v)]; }
};
/**
 * @brief Lowest Common Ancestor(最小共通祖先)
 * @docs docs/LCA.md
 * @see https://ei1333.github.io/luzhiled/snippets/tree/doubling-lowest-common-ancestor.html
 */
#line 4 "test/lca.test.cpp"
int main() {
    std::cin.tie(nullptr) -> sync_with_stdio(false);
    int n, q;
    std::cin >> n >> q;
    graph<false> g(n, 0);
    for(int i = 1; i < n; ++i) {
        int p;
        std::cin >> p;
        g.add(p, i);
    }
    LowestCommonAncestor lca(g);
    while(q--) {
        int u, v;
        std::cin >> u >> v;
        std::cout << lca.query(u, v) << '\n';
    }
}