test/LCA.test.cpp

• View this file on GitHub
• Last update: 2021-03-22 17:28:05+09:00
• Problem: https://judge.yosupo.jp/problem/lca

Depends on

• Graph/GraphTemplate.cpp
• Graph/LCA.cpp

Code

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

int main() {
	cin.tie(nullptr);
	ios_base::sync_with_stdio(false);
	int n, q;
	cin >> n >> q;
	Graph g(n);
	for (int i = 1; i < n; ++i) {
		int p;
		cin >> p;
		g[i].emplace_back(p, 1);
		g[p].emplace_back(i, 1);
	}

	LCA lca(g, 0);
	while (q--) {
		int u, v;
		cin >> u >> v;
		cout << lca(u, v) << '\n';
	}
}

#line 1 "test/LCA.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/lca"
#line 2 "Graph/GraphTemplate.cpp"
#include <vector>
#include <utility>
#include <iostream>
#include <limits>

using Weight = long long;
constexpr Weight INF = std::numeric_limits<Weight>::max();
struct Edge {
	int to;
	Weight cost;
	Edge() : to(-1), cost(-1) {}
	Edge(int _to, Weight _cost = 1) : to(_to), cost(_cost) {}
	friend bool operator<(const Edge& e1, const Edge& e2) {
		return e1.cost < e2.cost;
	}
	friend bool operator>(const Edge& e1, const Edge& e2) {
		return e1.cost > e2.cost;
	}
	friend std::ostream& operator<<(std::ostream& os, const Edge& e) {
		return os << "->" << e.to << '(' << e.cost << ')';
	}
};
using UnWeightedGraph = std::vector<std::vector<int>>;
using Graph = std::vector<std::vector<Edge>>;
struct Edge2 {
	int from, to;
	Weight cost;
	Edge2() : from(-1), to(-1), cost(0) {}
	Edge2(int _from, int _to, Weight _cost) : from(_from), to(_to), cost(_cost) {}
	friend bool operator<(const Edge2& e1, const Edge2& e2) {
		return e1.cost < e2.cost;
	}
	friend bool operator>(const Edge2& e1, const Edge2& e2) {
		return e1.cost > e2.cost;
	}
	friend std::ostream& operator<<(std::ostream& os, const Edge2& e) {
		return os << e.from << "->" << e.to << '(' << e.cost << ')';
	}
};
using UnWeightedEdges = std::vector<std::pair<int, int>>;
using Edges = std::vector<Edge2>;
using Matrix = std::vector<std::vector<Weight>>;

auto add_edge(UnWeightedGraph& graph, int v, int u) {
	graph[v].push_back(u);
	graph[u].push_back(v);
}
auto add_edge(Graph& graph, int v, int u, Weight cost) {
	graph[v].emplace_back(u, cost);
	graph[u].emplace_back(v, cost);
}
auto to_graph(const UnWeightedGraph& graph, Weight cost = 1) {
	Graph result(graph.size());
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (int v : graph[i]) {
			result[i].emplace_back(v, cost);
		}
	}
	return result;
}
auto to_unweighted_graph(const Graph& graph) {
	UnWeightedGraph result(graph.size());
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (auto [v, cost] : graph[i]) {
			result[i].push_back(v);
		}
	}
	return result;
}
auto to_edges(const UnWeightedGraph& graph, bool unique = false) {
	std::vector<std::pair<int, int>> edges;
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (int v : graph[i]) {
			if (!unique || static_cast<int>(i) < v) edges.emplace_back(i, v);
		}
	}
	return edges;
}
auto to_edges(const Graph& graph) {
	Edges edges;
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (auto [v, cost] : graph[i]) {
			edges.emplace_back(i, v, cost);
		}
	}
	return edges;
}
#line 5 "Graph/LCA.cpp"
#include <cmath>

class LCA {
	const std::size_t n;
	const int LOG;
	std::vector<int> dist;
	std::vector<std::vector<int>> table;
	void dfs(const Graph& graph, int v, int p, int d) {
		table[0][v] = p;
		dist[v] = d;
		for (auto e : graph[v])
			if (e.to != p) {
				dfs(graph, e.to, v, d + 1);
			}
	}

public:
	LCA(const Graph& graph, const int root)
	    : n(graph.size()), LOG(log2(n) + 1), dist(n), table(LOG, std::vector<int>(n)) {
		dfs(graph, root, -1, 0);
		for (int k = 0; k + 1 < LOG; ++k) {
			for (std::size_t v = 0; v < n; ++v) {
				if (table[k][v] < 0) {
					table[k + 1][v] = -1;
				} else {
					table[k + 1][v] = table[k][table[k][v]];
				}
			}
		}
	}
	int operator()(int u, int v) {
		if (dist[u] > dist[v]) std::swap(u, v);
		for (int k = 0; k < LOG; ++k) {
			if ((dist[v] - dist[u]) >> k & 1) {
				v = table[k][v];
			}
		}
		if (u == v) return u;
		for (int k = LOG - 1; k >= 0; --k) {
			if (table[k][u] != table[k][v]) {
				u = table[k][u];
				v = table[k][v];
			}
		}
		return table[0][u];
	}
};
#line 4 "test/LCA.test.cpp"
using namespace std;

int main() {
	cin.tie(nullptr);
	ios_base::sync_with_stdio(false);
	int n, q;
	cin >> n >> q;
	Graph g(n);
	for (int i = 1; i < n; ++i) {
		int p;
		cin >> p;
		g[i].emplace_back(p, 1);
		g[p].emplace_back(i, 1);
	}

	LCA lca(g, 0);
	while (q--) {
		int u, v;
		cin >> u >> v;
		cout << lca(u, v) << '\n';
	}
}

Back to top page