重心分解

tree/centroid_decomposition.hpp · GitHub で見る

共通ページ

#include "lib/tree/centroid_decomposition.hpp"

CentroidDecomposition<> tree(3);
tree.add_edge(0, 1);
tree.add_edge(1, 2);
tree.build();
int root = tree.root;
int distance = tree.dist(0, 2);
tree.for_each_ancestor(0, [&](int c, int d, int branch) {
	// c までの距離が d。branch は重複の除外に使う。
});

操作	内容・返り値	計算量
CentroidDecomposition<T = int>(n = 0)	頂点数と距離の型を指定	\(O(n)\)
add_edge(u,v,weight = 1)	重み付き無向辺を追加し、辺 ID を返す	償却 \(O(1)\)
build(distances = true)	重心分解を構築。false なら距離キャッシュを省く	\(O(n\log n)\)
has_distances()	構築済みの距離キャッシュを持つか	\(O(1)\)
ancestors(v)	重心分解木の根から v までの span<const Ancestor>	\(O(1)\)
for_each_ancestor(v,f)	v から根の順に f(c,d,branch) を呼ぶ	\(O(\log n)\)
is_ancestor(c,v)	重心分解木で c が v の祖先か。自身も含む	\(O(1)\)
subtree(c)	重心分解木の部分木に対応する order の区間 [in[c],out[c])	\(O(1)\)
centroid_lca(u,v)	重心分解木の最小共通祖先	\(O(\log n)\)
distance_to_ancestor(v,c)	元の木での v と重心祖先 c の距離	\(O(1)\)
dist(u,v)	元の木の経路上の辺重みの和	\(O(\log n)\)
Ancestor{centroid,branch,distance}	重心祖先／その祖先から v に向かう重心分解木の子／元の木での距離	\(O(1)\)
Edge{u,v,weight}	辺の両端と重み	\(O(1)\)
root	重心分解木の根。空の木では -1	\(O(1)\)
parent[v] / depth[v] / size[v]	重心分解木の親（根は -1）／深さ／部分木の頂点数	\(O(1)\)
in[v] / out[v] / order[i]	重心分解木の部分木の左端／右端／DFS 順の頂点	\(O(1)\)
n / edges[id]	全頂点数／入力した辺	\(O(1)\)

入力は頂点番号が 0 始まりの無向木です。空の木も扱えます。根は自動で決まり、辺重みには依存しません。size[c] は c を重心として選んだ時点の成分の頂点数でもあります。各走査は再帰を使いません。

T は辺重みと経路和を表せる型を指定します。重み付き木では CentroidDecomposition<long long> などを使います。負の重みも扱えますが、最近傍探索など三角不等式を使う応用では非負の重みが必要です。

Ancestor::branch は centroid から対象頂点に向かう重心分解木の子の頂点番号で、自身のレコードでは -1 です。同じ重心の同じ枝に属する頂点を除くことで、距離ごとの総和などの重複を取り除けます。枝の距離も、その親の重心から測ります。

例えば非負重みの木で、各重心について有効な頂点までの最小距離 best[c] を保存すれば、頂点の追加と最近傍クエリを書けます。次は最初の例の tree に対し、頂点 0 を有効にして頂点 2 からの最短距離を求めます。inf は距離和より十分大きく、加算がオーバーフローしない値にします。

const int inf = 1000000000;
std::vector<int> best(tree.n, inf);
int v = 0;
tree.for_each_ancestor(v, [&](int c, int d, int) {
	best[c] = std::min(best[c], d);
});
v = 2;
int answer = inf;
tree.for_each_ancestor(v, [&](int c, int d, int) {
	if (best[c] != inf) answer = std::min(answer, best[c] + d);
});
// answer = 2

build(false) でも重心分解木と部分木の区間は取得できます。ancestors・for_each_ancestor・distance_to_ancestor・dist には build(true) が必要です。ancestors の返す参照は再構築・辺の追加で無効になります。構築後の公開配列は参照用です。元の木の LCA が必要なら LCA を使います。

構築は \(O(n\log n)\) 時間、距離キャッシュを持つ場合は \(O(n\log n)\) 空間、持たない場合は \(O(n)\) 空間です。add_edge は償却 \(O(1)\)、for_each_ancestor・centroid_lca・dist は \(O(\log n)\)、それ以外のクエリは \(O(1)\) です。

検証コード

• test/centroid_decomposition_distance_frequency.test.cpp

コード

#ifndef CENTROID_DECOMPOSITION_HPP
#define CENTROID_DECOMPOSITION_HPP

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <limits>
#include <span>
#include <utility>
#include <vector>

template<class T = int> class CentroidDecomposition {
	bool built_ = false, distances_ = false;
	std::vector<std::size_t> offsets_;
	void check([[maybe_unused]] int v) const { assert(built_ && 0 <= v && v < n); }

  public:
	struct Edge { int u, v; T weight; };
	struct Ancestor { int centroid, branch; T distance; };

  private:
	std::vector<Ancestor> paths_;

  public:
	int n, root = -1;
	std::vector<Edge> edges;
	std::vector<int> parent, depth, size, in, out, order;
	explicit CentroidDecomposition(int n = 0) : n(n) {
		assert(0 <= n && n <= std::numeric_limits<int>::max() / 2);
		edges.reserve(std::max(0, n - 1));
	}
	int add_edge(int u, int v, T weight = T(1)) {
		assert(0 <= u && u < n && 0 <= v && v < n && u != v);
		built_ = distances_ = false;
		edges.push_back({u, v, weight});
		return int(edges.size()) - 1;
	}
	void build(bool distances = true) {
		assert(n >= 0 && edges.size() == std::size_t(std::max(0, n - 1)));
		built_ = distances_ = false;
		root = -1;
		parent.assign(n, -1), depth.assign(n, -1), size.assign(n, 1);
		in.resize(n), out.resize(n), order.clear();
		order.reserve(n);
		if (!distances) {
			std::vector<Ancestor>().swap(paths_);
			std::vector<std::size_t>().swap(offsets_);
		}
		if (!n) {
			paths_.clear(), offsets_.assign(1, 0);
			built_ = true, distances_ = distances;
			return;
		}
		std::vector<int> start(n + 1), adj(2 * edges.size());
		std::vector<T> weights(distances ? adj.size() : 0);
		for (const auto& e : edges) ++start[e.u + 1], ++start[e.v + 1];
		for (int v = 0; v < n; ++v) start[v + 1] += start[v];
		{
			auto cursor = start;
			for (const auto& e : edges) {
				int a = cursor[e.u]++, b = cursor[e.v]++;
				adj[a] = e.v, adj[b] = e.u;
				if (distances) weights[a] = weights[b] = e.weight;
			}
		}
		std::vector<int> work;
		work.reserve(n);
		work.push_back(0), parent[0] = -2;
		for (std::size_t i = 0; i < work.size(); ++i) {
			int v = work[i];
			for (int j = start[v]; j < start[v + 1]; ++j) {
				int u = adj[j];
				if (u == parent[v]) continue;
				assert(parent[u] == -1);
				parent[u] = v, work.push_back(u);
			}
		}
		assert(work.size() == std::size_t(n));
		for (std::size_t i = work.size(); --i;) size[parent[work[i]]] += size[work[i]];
		work.clear(), work.push_back(0), parent[0] = -1;
		while (!work.empty()) {
			int c = work.back(), p = parent[c], count = size[c];
			work.pop_back();
			while (count > 2) {
				int next = -1;
				for (int j = start[c]; j < start[c + 1]; ++j) {
					int u = adj[j];
					if (depth[u] < 0 && size[u] > count / 2) { next = u; break; }
				}
				if (next < 0) break;
				size[c] -= size[next], size[next] = count, c = next;
			}
			parent[c] = p, depth[c] = p < 0 ? 0 : depth[p] + 1;
			in[c] = int(order.size()), out[c] = in[c] + count;
			order.push_back(c);
			if (p < 0) root = c;
			if (count == 1) continue;
			for (int j = start[c + 1]; j-- > start[c];) {
				int u = adj[j];
				if (depth[u] < 0) parent[u] = c, work.push_back(u);
			}
		}
		if (distances) {
			offsets_.resize(n + 1), offsets_[0] = 0;
			for (int v = 0; v < n; ++v) offsets_[v + 1] = offsets_[v] + depth[v] + 1;
			paths_.resize(offsets_.back());
			std::vector<T> distance(n);
			std::vector<int> previous(n);
			for (int c : order) {
				work.clear(), work.push_back(c), previous[c] = -1, distance[c] = T(0);
				for (std::size_t i = 0; i < work.size(); ++i) {
					int v = work[i];
					paths_[offsets_[v] + depth[c]] = {c, -1, distance[v]};
					for (int j = start[v]; j < start[v + 1]; ++j) {
						int u = adj[j];
						if (u == previous[v] || depth[u] < depth[c]) continue;
						previous[u] = v, distance[u] = distance[v] + weights[j];
						work.push_back(u);
					}
				}
			}
			for (int v = 0; v < n; ++v)
				for (std::size_t i = offsets_[v]; i + 1 < offsets_[v + 1]; ++i)
					paths_[i].branch = paths_[i + 1].centroid;
		}
		built_ = true, distances_ = distances;
	}
	bool has_distances() const { return built_ && distances_; }
	std::span<const Ancestor> ancestors(int v) const {
		check(v);
		assert(distances_);
		return std::span<const Ancestor>(paths_).subspan(offsets_[v], offsets_[v + 1] - offsets_[v]);
	}
	template<class F> void for_each_ancestor(int v, F&& f) const {
		auto path = ancestors(v);
		for (std::size_t i = path.size(); i-- > 0;) {
			const auto& a = path[i];
			f(a.centroid, a.distance, a.branch);
		}
	}
	bool is_ancestor(int c, int v) const {
		check(c), check(v);
		return in[c] <= in[v] && in[v] < out[c];
	}
	std::pair<int, int> subtree(int c) const {
		check(c);
		return {in[c], out[c]};
	}
	int centroid_lca(int u, int v) const {
		check(u), check(v);
		while (depth[u] > depth[v]) u = parent[u];
		while (depth[v] > depth[u]) v = parent[v];
		while (u != v) u = parent[u], v = parent[v];
		return u;
	}
	T distance_to_ancestor(int v, int c) const {
		assert(distances_ && is_ancestor(c, v));
		check(v), check(c);
		return paths_[offsets_[v] + depth[c]].distance;
	}
	T dist(int u, int v) const {
		int c = centroid_lca(u, v);
		return distance_to_ancestor(u, c) + distance_to_ancestor(v, c);
	}
};

#endif

Back to top page