DSU on Tree

tree/dsu_on_tree.hpp · GitHub で見る

共通ページ

各頂点を根とする部分木について、色の種類数・最頻値などをまとめて求めます。最大の子の集計を残し、それ以外の子だけを追加し直します。Union-Find を使うライブラリではありません。

次は各部分木に現れる色の種類数を求める例です。color[v] は 0 以上 colors 未満とします。

#include "lib/tree/dsu_on_tree.hpp"

DSUOnTree tree(5);
tree.add_edge(0, 1);
tree.add_edge(0, 2);
tree.add_edge(1, 3);
tree.add_edge(1, 4);
tree.build(0);

std::vector<int> color{0, 1, 0, 2, 1};
int colors = 3, distinct = 0;
std::vector<int> count(colors), answer(tree.size());
tree.run(
	[&](int v) { distinct += count[color[v]]++ == 0; },
	[&](int v) { distinct -= --count[color[v]] == 0; },
	[&](int v) { answer[v] = distinct; }
);
// answer は 3, 2, 1, 1, 1。distinct と count は空の状態に戻る。

操作	内容・返り値	計算量
`DSUOnTree(n = 0)`	頂点数を指定。頂点番号は 0 始まり	\(O(n)\)
`size()` / `empty()` / `edge_count()`	頂点数／空判定／追加済み辺数	\(O(1)\)
`add_edge(u,v)`	無向辺を追加し、0 始まりの辺 ID を返す	償却 \(O(1)\)
`edge(id)`	追加した辺の両端の組への const 参照	\(O(1)\)
`build(root = 0)`	根を指定して前処理・再構築	\(O(n)\)
`run(add,remove,answer,keep = false)`	頂点ごとの追加・削除で部分木を集計	\(O(n\log n)\)
`run_clear(add,answer,clear,keep = false)`	個別削除の代わりに `clear()` で状態全体を空にする	\(O(n\log n)\) ＋回答・リセットの合計時間
`root()` / `parent(v)`	根／親。空の木の根と、根の親は -1	\(O(1)\)
`heavy_child(v)`	部分木が最大の子。葉なら -1。同数のときの選択は未規定	\(O(1)\)
`subtree_size(v)`	自身を含む部分木の頂点数	\(O(1)\)
`subtree(v)`	部分木に対応する `order()` の半開区間 `[l,r)`	\(O(1)\)
`vertex(i)` / `order()`	DFS 順の `i` 番目の頂点／全頂点の `span<const int>`	\(O(1)\)
`vertices(v)`	部分木の頂点列の `span<const int>`	\(O(1)\)

入力は連結な無向木で、辺をちょうど n-1 本追加してから build() を呼びます。空の木も扱え、その場合の実行ではコールバックを呼びません。根を変える場合は build(root) を呼び直します。辺を追加すると前処理は無効になります。返された span は、再構築・オブジェクトの破棄まで有効です。

add(v) は頂点を状態に追加し、remove(v) は追加済みの頂点を状態から除きます。answer(v) が呼ばれる瞬間には、状態に含まれる頂点が v の部分木と完全に一致します。各頂点について一度ずつ呼びますが、頂点番号順ではありません。答えは例のように頂点番号で保存してください。追加・削除の順序は未規定です。コールバックから木の再構築や状態を変える再入実行はしないでください。

実行開始時は状態を空にしてください。既定では実行終了時も空に戻るので、異なる集計を同じ木で繰り返せます。keep=true なら最後の根の集計を消さず、全頂点を追加した状態で終了します。再実行するときは利用者がその状態を空に戻してください。

個別に値を取り消しにくい集計には run_clear を使えます。clear() は状態全体を空にする処理です。次は最初の例の tree・color・colors を使い、各部分木の色の最大出現回数を求めます。触れた色だけをリセットすることで、毎回配列全体を初期化する二乗時間の処理を避けます。

std::vector<int> count2(colors), touched, mode(tree.size());
int maximum = 0;
tree.run_clear(
	[&](int v) {
		int c = color[v];
		if (count2[c] == 0) touched.push_back(c);
		maximum = std::max(maximum, ++count2[c]);
	},
	[&](int v) { mode[v] = maximum; },
	[&]() {
		for (int c : touched) count2[c] = 0;
		touched.clear();
		maximum = 0;
	}
);

各走査は再帰を使わず、一本道の木でもコールスタックを消費しません。重い子を先頭にした DFS 順を逆から走査し、軽い子の追加を連続した区間で行います。頂点ごとのイベント列や仮想関数は使いません。

前処理は時間・空間とも \(O(n)\)。add・remove がそれぞれ \(O(1)\) なら run は \(O(n\log n)\)、answer の呼び出しはちょうど \(n\) 回です。run_clear の追加回数も \(O(n\log n)\)、リセット回数は高々 \(n\) 回で、実行時間には answer と clear の合計時間を加えます。各 clear が消す頂点数以下の時間なら全体で \(O(n\log n)\) です。走査自体の追加領域は \(O(1)\)、ライブラリ全体の空間は \(O(n)\) です。

依存ライブラリ

無向グラフの CSR

検証コード

test/dsu_on_tree_subtree_sum.test.cpp

コード

#line 1 "tree/dsu_on_tree.hpp"

#line 1 "components/undirected_csr.hpp"

#include <cassert>
#include <cstddef>
#include <limits>
#include <memory>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>

namespace undirected_csr_detail {
template <class Arc>
void build(int n, std::span<const std::pair<int, int>> edges, std::vector<int>& start, std::unique_ptr<Arc[]>& adj) {
	assert(0 <= n && n < std::numeric_limits<int>::max());
	assert(edges.size() <= std::size_t(std::numeric_limits<int>::max() / 2));
	start.assign(std::size_t(n) + 1, 0);
	for (auto [u, v] : edges) {
		assert(0 <= u && u < n && 0 <= v && v < n);
		++start[u + 1], ++start[v + 1];
	}
	for (int v = 0; v < n; ++v) start[v + 1] += start[v];
	adj.reset(new Arc[2 * edges.size()]);
	std::vector<int> cursor(start.begin(), start.end() - 1);
	for (int i = 0; i < static_cast<int>(edges.size()); ++i) {
		auto [u, v] = edges[i];
		if constexpr (std::is_same_v<Arc, int>) adj[cursor[u]++] = v, adj[cursor[v]++] = u;
		else adj[cursor[u]++] = {v, i}, adj[cursor[v]++] = {u, i};
	}
}
}

struct UndirectedCSR {
	struct Arc { int to, id; };
	int n;
	std::vector<int> start;
	std::unique_ptr<Arc[]> adj;
	UndirectedCSR(int n, std::span<const std::pair<int, int>> edges) : n(n) {
		undirected_csr_detail::build(n, edges, start, adj);
	}
	std::span<const Arc> operator[](int v) const {
		assert(0 <= v && v < n);
		return {adj.get() + start[v], std::size_t(start[v + 1] - start[v])};
	}
};

struct UndirectedVertexCSR {
	int n;
	std::vector<int> start;
	std::unique_ptr<int[]> adj;
	UndirectedVertexCSR(int n, std::span<const std::pair<int, int>> edges) : n(n) {
		undirected_csr_detail::build(n, edges, start, adj);
	}
	std::span<const int> operator[](int v) const {
		assert(0 <= v && v < n);
		return {adj.get() + start[v], std::size_t(start[v + 1] - start[v])};
	}
};

#line 5 "tree/dsu_on_tree.hpp"
#include <algorithm>
#line 11 "tree/dsu_on_tree.hpp"

class DSUOnTree {
	int n_, root_ = -1;
	bool built_ = false;
	std::vector<std::pair<int, int>> edges_;
	std::vector<int> parent_, heavy_, in_, out_, order_;
	void check([[maybe_unused]] int v) const { assert(built_ && 0 <= v && v < n_); }
	template<class Add, class Answer, class Clear>
	void execute(Add& add, Answer& answer, Clear& clear, bool keep) const {
		assert(built_);
		for (int i = n_; i-- > 0;) {
			int v = order_[i], h = heavy_[v];
			for (int j = h == -1 ? i + 1 : out_[h]; j < out_[v]; ++j) add(order_[j]);
			add(v);
			answer(v);
			if (v == root_ ? !keep : heavy_[parent_[v]] != v) clear(v);
		}
	}

public:
	explicit DSUOnTree(int n = 0) : n_(n) {
		assert(0 <= n && n <= std::numeric_limits<int>::max() / 2);
		edges_.reserve(std::max(0, n - 1));
	}
	int size() const { return n_; }
	bool empty() const { return n_ == 0; }
	int edge_count() const { return int(edges_.size()); }
	int add_edge(int u, int v) {
		assert(0 <= u && u < n_ && 0 <= v && v < n_ && u != v);
		assert(edges_.size() < std::size_t(std::max(0, n_ - 1)));
		built_ = false;
		edges_.emplace_back(u, v);
		return edge_count() - 1;
	}
	const std::pair<int, int>& edge(int id) const {
		assert(0 <= id && id < edge_count());
		return edges_[id];
	}
	void build(int root = 0) {
		assert(edge_count() == std::max(0, n_ - 1));
		assert(!n_ || (0 <= root && root < n_));
		built_ = false;
		root_ = n_ ? root : -1;
		parent_.assign(n_, -2);
		heavy_.assign(n_, -1);
		in_.resize(n_);
		out_.assign(n_, 1);
		order_.clear();
		if (!n_) { built_ = true; return; }
		UndirectedCSR graph(n_, edges_);
		order_.reserve(n_);
		order_.push_back(root);
		parent_[root] = -1;
		for (std::size_t i = 0; i < order_.size(); ++i) {
			int v = order_[i];
			for (auto arc : graph[v]) {
				int to = arc.to;
				if (to == parent_[v]) continue;
				assert(parent_[to] == -2);
				if (parent_[to] != -2) continue;
				parent_[to] = v;
				order_.push_back(to);
			}
		}
		assert(int(order_.size()) == n_);
		for (int i = n_ - 1; i > 0; --i) {
			int v = order_[i], p = parent_[v];
			out_[p] += out_[v];
			if (heavy_[p] == -1 || out_[v] > out_[heavy_[p]]) heavy_[p] = v;
		}
		int timer = 0, top = n_ - 1;
		order_[top] = root;
		while (top < n_) {
			int v = order_[top++];
			in_[v] = timer;
			order_[timer++] = v;
			for (auto arc : graph[v])
				if (parent_[arc.to] == v && arc.to != heavy_[v]) order_[--top] = arc.to;
			if (heavy_[v] != -1) order_[--top] = heavy_[v];
		}
		for (int v = 0; v < n_; ++v) out_[v] += in_[v];
		built_ = true;
	}
	int root() const { assert(built_); return root_; }
	int parent(int v) const { check(v); return parent_[v]; }
	int heavy_child(int v) const { check(v); return heavy_[v]; }
	int subtree_size(int v) const { check(v); return out_[v] - in_[v]; }
	std::pair<int, int> subtree(int v) const { check(v); return {in_[v], out_[v]}; }
	int vertex(int i) const {
		assert(built_ && 0 <= i && i < n_);
		return order_[i];
	}
	std::span<const int> order() const { assert(built_); return order_; }
	std::span<const int> vertices(int v) const {
		check(v);
		return std::span<const int>(order_).subspan(in_[v], out_[v] - in_[v]);
	}
	template<class Add, class Remove, class Answer>
	void run(Add&& add, Remove&& remove, Answer&& answer, bool keep = false) const {
		auto clear = [&](int v) {
			for (int i = in_[v]; i < out_[v]; ++i) remove(order_[i]);
		};
		execute(add, answer, clear, keep);
	}
	template<class Add, class Answer, class Clear>
	void run_clear(Add&& add, Answer&& answer, Clear&& clear, bool keep = false) const {
		auto reset = [&](int) { clear(); };
		execute(add, answer, reset, keep);
	}
};

#ifndef DSU_ON_TREE_HPP
#define DSU_ON_TREE_HPP

#include "../components/undirected_csr.hpp"
#include <algorithm>
#include <cassert>
#include <limits>
#include <span>
#include <utility>
#include <vector>

class DSUOnTree {
	int n_, root_ = -1;
	bool built_ = false;
	std::vector<std::pair<int, int>> edges_;
	std::vector<int> parent_, heavy_, in_, out_, order_;
	void check([[maybe_unused]] int v) const { assert(built_ && 0 <= v && v < n_); }
	template<class Add, class Answer, class Clear>
	void execute(Add& add, Answer& answer, Clear& clear, bool keep) const {
		assert(built_);
		for (int i = n_; i-- > 0;) {
			int v = order_[i], h = heavy_[v];
			for (int j = h == -1 ? i + 1 : out_[h]; j < out_[v]; ++j) add(order_[j]);
			add(v);
			answer(v);
			if (v == root_ ? !keep : heavy_[parent_[v]] != v) clear(v);
		}
	}

  public:
	explicit DSUOnTree(int n = 0) : n_(n) {
		assert(0 <= n && n <= std::numeric_limits<int>::max() / 2);
		edges_.reserve(std::max(0, n - 1));
	}
	int size() const { return n_; }
	bool empty() const { return n_ == 0; }
	int edge_count() const { return int(edges_.size()); }
	int add_edge(int u, int v) {
		assert(0 <= u && u < n_ && 0 <= v && v < n_ && u != v);
		assert(edges_.size() < std::size_t(std::max(0, n_ - 1)));
		built_ = false;
		edges_.emplace_back(u, v);
		return edge_count() - 1;
	}
	const std::pair<int, int>& edge(int id) const {
		assert(0 <= id && id < edge_count());
		return edges_[id];
	}
	void build(int root = 0) {
		assert(edge_count() == std::max(0, n_ - 1));
		assert(!n_ || (0 <= root && root < n_));
		built_ = false;
		root_ = n_ ? root : -1;
		parent_.assign(n_, -2);
		heavy_.assign(n_, -1);
		in_.resize(n_);
		out_.assign(n_, 1);
		order_.clear();
		if (!n_) { built_ = true; return; }
		UndirectedCSR graph(n_, edges_);
		order_.reserve(n_);
		order_.push_back(root);
		parent_[root] = -1;
		for (std::size_t i = 0; i < order_.size(); ++i) {
			int v = order_[i];
			for (auto arc : graph[v]) {
				int to = arc.to;
				if (to == parent_[v]) continue;
				assert(parent_[to] == -2);
				if (parent_[to] != -2) continue;
				parent_[to] = v;
				order_.push_back(to);
			}
		}
		assert(int(order_.size()) == n_);
		for (int i = n_ - 1; i > 0; --i) {
			int v = order_[i], p = parent_[v];
			out_[p] += out_[v];
			if (heavy_[p] == -1 || out_[v] > out_[heavy_[p]]) heavy_[p] = v;
		}
		int timer = 0, top = n_ - 1;
		order_[top] = root;
		while (top < n_) {
			int v = order_[top++];
			in_[v] = timer;
			order_[timer++] = v;
			for (auto arc : graph[v])
				if (parent_[arc.to] == v && arc.to != heavy_[v]) order_[--top] = arc.to;
			if (heavy_[v] != -1) order_[--top] = heavy_[v];
		}
		for (int v = 0; v < n_; ++v) out_[v] += in_[v];
		built_ = true;
	}
	int root() const { assert(built_); return root_; }
	int parent(int v) const { check(v); return parent_[v]; }
	int heavy_child(int v) const { check(v); return heavy_[v]; }
	int subtree_size(int v) const { check(v); return out_[v] - in_[v]; }
	std::pair<int, int> subtree(int v) const { check(v); return {in_[v], out_[v]}; }
	int vertex(int i) const {
		assert(built_ && 0 <= i && i < n_);
		return order_[i];
	}
	std::span<const int> order() const { assert(built_); return order_; }
	std::span<const int> vertices(int v) const {
		check(v);
		return std::span<const int>(order_).subspan(in_[v], out_[v] - in_[v]);
	}
	template<class Add, class Remove, class Answer>
	void run(Add&& add, Remove&& remove, Answer&& answer, bool keep = false) const {
		auto clear = [&](int v) {
			for (int i = in_[v]; i < out_[v]; ++i) remove(order_[i]);
		};
		execute(add, answer, clear, keep);
	}
	template<class Add, class Answer, class Clear>
	void run_clear(Add&& add, Answer&& answer, Clear&& clear, bool keep = false) const {
		auto reset = [&](int) { clear(); };
		execute(add, answer, reset, keep);
	}
};

#endif

#line 1 "tree/dsu_on_tree.hpp"



#line 1 "components/undirected_csr.hpp"



#include <cassert>
#include <cstddef>
#include <limits>
#include <memory>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>

namespace undirected_csr_detail {
template <class Arc>
void build(int n, std::span<const std::pair<int, int>> edges, std::vector<int>& start, std::unique_ptr<Arc[]>& adj) {
	assert(0 <= n && n < std::numeric_limits<int>::max());
	assert(edges.size() <= std::size_t(std::numeric_limits<int>::max() / 2));
	start.assign(std::size_t(n) + 1, 0);
	for (auto [u, v] : edges) {
		assert(0 <= u && u < n && 0 <= v && v < n);
		++start[u + 1], ++start[v + 1];
	}
	for (int v = 0; v < n; ++v) start[v + 1] += start[v];
	adj.reset(new Arc[2 * edges.size()]);
	std::vector<int> cursor(start.begin(), start.end() - 1);
	for (int i = 0; i < static_cast<int>(edges.size()); ++i) {
		auto [u, v] = edges[i];
		if constexpr (std::is_same_v<Arc, int>) adj[cursor[u]++] = v, adj[cursor[v]++] = u;
		else adj[cursor[u]++] = {v, i}, adj[cursor[v]++] = {u, i};
	}
}
}

struct UndirectedCSR {
	struct Arc { int to, id; };
	int n;
	std::vector<int> start;
	std::unique_ptr<Arc[]> adj;
	UndirectedCSR(int n, std::span<const std::pair<int, int>> edges) : n(n) {
		undirected_csr_detail::build(n, edges, start, adj);
	}
	std::span<const Arc> operator[](int v) const {
		assert(0 <= v && v < n);
		return {adj.get() + start[v], std::size_t(start[v + 1] - start[v])};
	}
};

struct UndirectedVertexCSR {
	int n;
	std::vector<int> start;
	std::unique_ptr<int[]> adj;
	UndirectedVertexCSR(int n, std::span<const std::pair<int, int>> edges) : n(n) {
		undirected_csr_detail::build(n, edges, start, adj);
	}
	std::span<const int> operator[](int v) const {
		assert(0 <= v && v < n);
		return {adj.get() + start[v], std::size_t(start[v + 1] - start[v])};
	}
};


#line 5 "tree/dsu_on_tree.hpp"
#include <algorithm>
#line 11 "tree/dsu_on_tree.hpp"

class DSUOnTree {
	int n_, root_ = -1;
	bool built_ = false;
	std::vector<std::pair<int, int>> edges_;
	std::vector<int> parent_, heavy_, in_, out_, order_;
	void check([[maybe_unused]] int v) const { assert(built_ && 0 <= v && v < n_); }
	template<class Add, class Answer, class Clear>
	void execute(Add& add, Answer& answer, Clear& clear, bool keep) const {
		assert(built_);
		for (int i = n_; i-- > 0;) {
			int v = order_[i], h = heavy_[v];
			for (int j = h == -1 ? i + 1 : out_[h]; j < out_[v]; ++j) add(order_[j]);
			add(v);
			answer(v);
			if (v == root_ ? !keep : heavy_[parent_[v]] != v) clear(v);
		}
	}

  public:
	explicit DSUOnTree(int n = 0) : n_(n) {
		assert(0 <= n && n <= std::numeric_limits<int>::max() / 2);
		edges_.reserve(std::max(0, n - 1));
	}
	int size() const { return n_; }
	bool empty() const { return n_ == 0; }
	int edge_count() const { return int(edges_.size()); }
	int add_edge(int u, int v) {
		assert(0 <= u && u < n_ && 0 <= v && v < n_ && u != v);
		assert(edges_.size() < std::size_t(std::max(0, n_ - 1)));
		built_ = false;
		edges_.emplace_back(u, v);
		return edge_count() - 1;
	}
	const std::pair<int, int>& edge(int id) const {
		assert(0 <= id && id < edge_count());
		return edges_[id];
	}
	void build(int root = 0) {
		assert(edge_count() == std::max(0, n_ - 1));
		assert(!n_ || (0 <= root && root < n_));
		built_ = false;
		root_ = n_ ? root : -1;
		parent_.assign(n_, -2);
		heavy_.assign(n_, -1);
		in_.resize(n_);
		out_.assign(n_, 1);
		order_.clear();
		if (!n_) { built_ = true; return; }
		UndirectedCSR graph(n_, edges_);
		order_.reserve(n_);
		order_.push_back(root);
		parent_[root] = -1;
		for (std::size_t i = 0; i < order_.size(); ++i) {
			int v = order_[i];
			for (auto arc : graph[v]) {
				int to = arc.to;
				if (to == parent_[v]) continue;
				assert(parent_[to] == -2);
				if (parent_[to] != -2) continue;
				parent_[to] = v;
				order_.push_back(to);
			}
		}
		assert(int(order_.size()) == n_);
		for (int i = n_ - 1; i > 0; --i) {
			int v = order_[i], p = parent_[v];
			out_[p] += out_[v];
			if (heavy_[p] == -1 || out_[v] > out_[heavy_[p]]) heavy_[p] = v;
		}
		int timer = 0, top = n_ - 1;
		order_[top] = root;
		while (top < n_) {
			int v = order_[top++];
			in_[v] = timer;
			order_[timer++] = v;
			for (auto arc : graph[v])
				if (parent_[arc.to] == v && arc.to != heavy_[v]) order_[--top] = arc.to;
			if (heavy_[v] != -1) order_[--top] = heavy_[v];
		}
		for (int v = 0; v < n_; ++v) out_[v] += in_[v];
		built_ = true;
	}
	int root() const { assert(built_); return root_; }
	int parent(int v) const { check(v); return parent_[v]; }
	int heavy_child(int v) const { check(v); return heavy_[v]; }
	int subtree_size(int v) const { check(v); return out_[v] - in_[v]; }
	std::pair<int, int> subtree(int v) const { check(v); return {in_[v], out_[v]}; }
	int vertex(int i) const {
		assert(built_ && 0 <= i && i < n_);
		return order_[i];
	}
	std::span<const int> order() const { assert(built_); return order_; }
	std::span<const int> vertices(int v) const {
		check(v);
		return std::span<const int>(order_).subspan(in_[v], out_[v] - in_[v]);
	}
	template<class Add, class Remove, class Answer>
	void run(Add&& add, Remove&& remove, Answer&& answer, bool keep = false) const {
		auto clear = [&](int v) {
			for (int i = in_[v]; i < out_[v]; ++i) remove(order_[i]);
		};
		execute(add, answer, clear, keep);
	}
	template<class Add, class Answer, class Clear>
	void run_clear(Add&& add, Answer&& answer, Clear&& clear, bool keep = false) const {
		auto reset = [&](int) { clear(); };
		execute(add, answer, reset, keep);
	}
};