7K4B blog

猫でも分かる何か

NyaaLIB::GT_NyaaLCA

LCAライブラリ
木の最近共通祖先 (Lowest Common Ancestor)
前処理  O(N \log N) することで以下のクエリに  O(\log N) で答える

  • Ancestor(x, y): 木頂点 (x, y) に対する最近共通祖先
  • Dist(x, y): 全頂点対間距離(任意の2頂点 (x, y) 間の距離)
  • IsOnPath(x, y, z): x-y パス上に頂点 z が存在するか判定
#include <bits/stdc++.h>

namespace NyaaLIB
{
	/**
	* LCAライブラリ (Lowest Common Ancestor)
	* 前処理O(N*log(N)), クエリO(log(N))
	* @note
	* Ancestor(x,y) 木頂点(x,y)に対する最近共通祖先
	* Dist(x,y) 木頂点(x,y)間の距離
	* IsOnPath(x,y,z) x-y パス上に頂点 z が存在するか判定
	**/

	struct GT_NyaaLCA
	{
		using ll = long long;
		using vll = std::vector<ll>;
		using vvll = std::vector<std::vector<ll>>;
		using tree = std::vector<std::vector<ll>>;
		vll d;
		vvll p;
		tree g;
		GT_NyaaLCA(tree& g) : d(ll(g.size())), g(g)
		{
			ll k = 1;
			while ((1ll << k) < ll(g.size())) k++;
			p.assign(k, vll(ll(g.size())));
			auto GTL_NyaaTreeDFS = [&](auto self, tree& g, ll now, ll root, ll depth) -> void
			{
				p[0][now] = root, d[now] = depth;
				for (auto& next : g[now])
				{
					if (next == root) continue;
					self(self, g, next, now, depth + 1);
				}
			};
			GTL_NyaaTreeDFS(GTL_NyaaTreeDFS, g, 0, -1, 0);
			for (ll i = 0; i + 1 < k; i++) for (ll j = 0; j < ll(g.size()); j++)
			{
				if (p[i][j] < 0) p[i + 1][j] = -1;
				else p[i + 1][j] = p[i][p[i][j]];
			}
		}
		ll Ancestor(ll x, ll y)
		{
			if (d[x] > d[y]) std::swap(x, y);
			for (ll i = 0; i < ll(p.size()); i++) if ((d[y] - d[x]) >> i & 1) y = p[i][y];
			if (x == y) return x;
			for (ll i = ll(p.size()) - 1; i >= 0; i--) if (p[i][x] != p[i][y]) x = p[i][x], y = p[i][y];
			return p[0][x];
		}
		ll Dist(ll x, ll y) { return d[x] + d[y] - 2 * d[Ancestor(x, y)]; }
		bool IsOnPath(ll x, ll y, ll z) { return Dist(x, z) + Dist(z, y) == Dist(x, y); }
	};
}

実装例
https://atcoder.jp/contests/abc220/submissions/26677587