Auxiliary Tree

tree/auxiliary_tree.hpp · GitHub で見る

共通ページ

大きな木から、指定した頂点と、それらをつなぐために必要な LCA だけを残した小さな木を作ります。通常は、元の木を build() してから AuxiliaryTree を作り、vt.build(vertices) に使いたい頂点番号を渡します。入力の順番を気にする必要はありません。

#include "lib/tree/lca.hpp"
#include "lib/tree/auxiliary_tree.hpp"
#include <iostream>

int main() {
	LCA tree(5);
	tree.add_edge(0, 1);
	tree.add_edge(1, 2);
	tree.add_edge(1, 3);
	tree.add_edge(3, 4);
	tree.build(0);

	AuxiliaryTree vt(tree);
	std::vector<int> vertices{2, 4};
	vt.build(vertices);
	for (int i = 0; i < vt.size(); ++i) {
		const auto &node = vt[i];
		if (node.parent == -1) continue;
		std::cout << node.vertex << ' ' << vt[node.parent].vertex
		          << ' ' << node.length << '\n';
	}
}

出力は 4 1 2、2 1 1 の 2 行です。各行は「元の頂点番号、補助木での親の元の頂点番号、親までの元の辺数」です。頂点 2 と 4 に加え、LCA の頂点 1 が残ります。元の木の頂点 3 は辺の途中にまとめられ、根だった頂点 0 も補助木には入りません。

LCA または HLD を先に build してください。その根に基づく LCA を使い、補助木の根は指定頂点全体の LCA になります。

操作	内容・返り値	計算量
`AuxiliaryTree(tree)`	前処理済みの木を参照して作成。型は推論可能	\(O(1)\)
`build(vertices)`	頂点列から構築し、`nodes()` と同じ配列への const 参照を返す	LCA: \(O(k)\)、HLD: \(O(k\log n)\)
`build_sorted(vertices)`	元の木の DFS 番号の小さい順に並んだ頂点列から構築。内部のソートを省く。通常は `build` を使う	LCA: \(O(k)\)、HLD: \(O(k\log n)\)
`size()`	補助木の頂点数	\(O(1)\)
`empty()`	補助木が空か	\(O(1)\)
`root()`	根の圧縮 ID。空なら -1、そうでなければ `size() - 1`	\(O(1)\)
`operator[](i)`	圧縮 ID `i` の `Node` への const 参照	\(O(1)\)
`nodes()`	全 `Node` の配列への const 参照	\(O(1)\)
`to_graph()`	圧縮 ID を頂点番号とする無向木の隣接リストを `vector<vector<int>>` で返す	\(O(m)\)

`Node` の要素	内容
`vertex`	元の木での頂点番号
`parent`	親の圧縮 ID。根は -1
`length`	親までの元の辺数。根は 0
`selected`	入力に含まれる頂点か。LCA として追加しただけの頂点は false

入力は std::vector<int>・std::array<int,k> などの std::span<const int> に変換できる列です。入力を書き換えず、重複は一つにまとめます。空入力にも対応し、異なる指定頂点が \(k\) 個なら補助木は最大 \(2k-1\) 頂点です。

元の頂点番号と、補助木の配列を参照するための圧縮 ID は別物です。vt[i].vertex は元の頂点番号、vt[i].parent は補助木の ID です。親の元の頂点番号を取得するには、根でないことを確認して vt[vt[i].parent].vertex とします。

圧縮 ID は子から親の順です。根以外は必ず i < vt[i].parent なので、ID の昇順で木 DP、降順で根からの伝播ができます。元の頂点番号や DFS の行きがけ順とは異なります。

std::vector<std::vector<int>> graph = vt.to_graph();

graph の頂点数は vt.size() で、graph[i] には圧縮 ID i の親・子の ID が入ります。元の頂点番号には vt[i].vertex、根には vt.root() を使います。辺の長さは隣接リストに含めず、子の length で取得します。空の補助木には空配列、1 頂点には空の隣接リスト 1 つを返します。返り値は独立した配列なので、書き換えや補助木の再構築の影響を互いに受けません。

build_sorted を使う場合の「DFS 番号」は、元の木で頂点を初めて訪れる順に 0, 1, 2, … と付けた番号です。tree.subtree(v) はその順序での部分木の区間 [l,r) を返し、.first、つまり l が頂点 v の DFS 番号です。頂点番号 v でソートすることとは異なります。

例えば、頂点 2 の DFS 番号が 4、頂点 4 の DFS 番号が 3 なら、build_sorted に渡せる順番は {4,2} です。任意の列をこの順番に直すコードは次のとおりです。上の例の vt.build(vertices) と置き換えられます。

std::sort(vertices.begin(), vertices.end(), [&](int u, int v) {
	return tree.subtree(u).first < tree.subtree(v).first;
});
vt.build_sorted(vertices);

通常の build は DFS 番号を整数キーとして基数ソートし、256件未満では比較ソートを使います。自分で上の比較ソートをすると \(O(k\log k)\) 時間かかります。build_sorted が役立つのは、木を DFS 順に走査しながら頂点を集めた場合など、入力が最初からこの順番になっているときです。順番が分からなければ build を使ってください。

重み付き距離は、元の木の根からの重み和を distance[v] として distance[vt[i].vertex] - distance[vt[vt[i].parent].vertex] で得られます。length は重みにかかわらず辺数です。

元の木は補助木より長く生存させ、移動・破棄しないでください。元の木の前処理をやり直した後は補助木も再構築してください。返された配列・要素への参照は、次の構築で内容が変わり、再確保で無効になります。作業配列の容量を再利用し、各構築で元の木の全頂点を走査しません。

Tree に必要な操作は lca(u,v)、部分木の DFS 区間を返す subtree(v)、辺数による深さ depth[v] です。

元の頂点数を \(n\)、入力列長を \(k\) とすると、LCA の前処理は時間・空間 \(O(n \log n)\)、補助木の構築は \(O(k)\) です。HLD は前処理・空間 \(O(n)\)、構築 \(O(k \log n)\) です。ソート済みの入力でも同じ計算量ですが、build_sorted はソートとその作業領域の確保を省けます。補助木の追加空間は最大の入力列長に対し \(O(k)\)、生成後の参照操作は \(O(1)\) です。to_graph() は補助木の頂点数を \(m\) として \(O(m)\) 時間・追加空間です。

依存ライブラリ

64 bit キーの安定基数ソート

検証コード

test/auxiliary_tree.test.cpp

コード

#line 1 "tree/auxiliary_tree.hpp"

#line 1 "components/radix_sort.hpp"

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>

template <int Bits, class T, class Key>
void stable_radix_sort(std::vector<T> &values, Key &&key) {
	static_assert(1 <= Bits && Bits <= 16);
	constexpr unsigned buckets = 1U << Bits, mask = buckets - 1;
	using K = std::remove_cvref_t<decltype(key(std::declval<const T &>()))>;
	static_assert(std::is_integral_v<K> && std::is_unsigned_v<K> && sizeof(K) <= sizeof(std::uint64_t));
	if (values.size() < 2) return;
	if (values.size() < 256) {
		std::stable_sort(values.begin(), values.end(), [&](const T &a, const T &b) { return key(a) < key(b); });
		return;
	}
	std::uint64_t first = key(values.front()), previous = first, varying = 0;
	bool sorted = true;
	for (std::size_t i = 1; i < values.size(); ++i) {
		std::uint64_t current = key(values[i]);
		varying |= first ^ current;
		sorted &= previous <= current;
		previous = current;
	}
	if (sorted) return;
	std::vector<T> scratch(values.size());
	using Counts = std::conditional_t<(Bits <= 11), std::array<std::size_t, buckets>, std::vector<std::size_t>>;
	Counts count = [] {
		if constexpr (Bits <= 11) return Counts{};
		else return Counts(buckets);
	}();
	for (int shift = 0; shift < 64; shift += Bits) {
		if (!((varying >> shift) & mask)) continue;
		std::fill(count.begin(), count.end(), 0);
		for (const auto &value : values) ++count[(static_cast<std::uint64_t>(key(value)) >> shift) & mask];
		std::size_t sum = 0;
		for (auto &x : count) {
			auto n = x;
			x = sum, sum += n;
		}
		for (auto &&value : values) scratch[count[(static_cast<std::uint64_t>(key(value)) >> shift) & mask]++] = std::move(value);
		values.swap(scratch);
	}
}

template <class T, class Key>
void stable_radix_sort(std::vector<T> &values, Key key) {
	stable_radix_sort<8>(values, std::move(key));
}

namespace radix_sort_detail {
template <int Bits = 8, class T, class Key>
void sort_by_key(std::vector<T> &values, Key key) {
	using K = std::remove_cvref_t<decltype(key(std::declval<const T &>()))>;
	if constexpr (std::is_integral_v<K> && !std::is_same_v<K, bool> && sizeof(K) <= 8 && std::is_default_constructible_v<T>) {
		using U = std::make_unsigned_t<K>;
		constexpr U sign = std::is_signed_v<K> ? U(1) << (sizeof(K) * 8 - 1) : U(0);
		stable_radix_sort<Bits>(values, [&](const T &value) { return U(U(key(value)) ^ sign); });
	} else std::sort(values.begin(), values.end(), [&](const T &a, const T &b) { return key(a) < key(b); });
}
}

#line 6 "tree/auxiliary_tree.hpp"
#include <cassert>
#include <span>
#line 10 "tree/auxiliary_tree.hpp"

template <class Tree> class AuxiliaryTree {
  public:
	struct Node {
		int vertex, parent, length;
		bool selected;
	};

private:
	struct Frame {
		int vertex, child;
		bool selected;
	};
	const Tree *tree_;
	std::vector<std::pair<int, int>> keys_;
	std::vector<Frame> stack_;
	std::vector<Node> nodes_;
	int finish() {
		Frame current = stack_.back();
		stack_.pop_back();
		int id = size();
		for (int child = current.child; child != -1;) {
			int next = nodes_[child].parent;
			nodes_[child].parent = id;
			nodes_[child].length = tree_->depth[nodes_[child].vertex] - tree_->depth[current.vertex];
			child = next;
		}
		nodes_.push_back({current.vertex, -1, 0, current.selected});
		return id;
	}
	void attach(int child) {
		nodes_[child].parent = stack_.back().child;
		stack_.back().child = child;
	}
	const std::vector<Node> &construct(std::span<const int> vertices, bool sorted) {
		keys_.clear();
		for (int v : vertices) keys_.emplace_back(tree_->subtree(v).first, v);
		if (!sorted) {
			if (keys_.size() < 256) std::sort(keys_.begin(), keys_.end());
			else stable_radix_sort(keys_, [](const auto &x) { return static_cast<unsigned>(x.first); });
		}
		else assert(std::is_sorted(keys_.begin(), keys_.end()));
		nodes_.clear();
		stack_.clear();
		int previous = -1;
		for (auto [position, v] : keys_) {
			(void)position;
			if (v == previous) continue;
			previous = v;
			if (!stack_.empty()) {
				int w = tree_->lca(stack_.back().vertex, v);
				while (tree_->depth[stack_.back().vertex] > tree_->depth[w]) {
					int child = finish();
					if (stack_.empty() || tree_->depth[stack_.back().vertex] < tree_->depth[w])
						stack_.push_back({w, -1, false});
					attach(child);
				}
			}
			stack_.push_back({v, -1, true});
		}
		while (stack_.size() > 1) attach(finish());
		if (!stack_.empty()) finish();
		return nodes_;
	}

public:
	explicit AuxiliaryTree(const Tree &tree) : tree_(&tree) {}
	AuxiliaryTree(const Tree &&) = delete;
	int size() const { return static_cast<int>(nodes_.size()); }
	bool empty() const { return nodes_.empty(); }
	int root() const { return size() - 1; }
	const Node &operator[](int i) const {
		assert(0 <= i && i < size());
		return nodes_[i];
	}
	const std::vector<Node> &nodes() const { return nodes_; }
	std::vector<std::vector<int>> to_graph() const {
		std::vector<std::vector<int>> graph(size());
		std::vector<int> degree(size());
		for (int i = 0; i < size(); ++i) if (nodes_[i].parent >= 0) {
			++degree[i];
			++degree[nodes_[i].parent];
		}
		for (int i = 0; i < size(); ++i) graph[i].reserve(degree[i]);
		for (int i = 0; i < size(); ++i) if (nodes_[i].parent >= 0) {
			int parent = nodes_[i].parent;
			graph[i].push_back(parent);
			graph[parent].push_back(i);
		}
		return graph;
	}
	const std::vector<Node> &build(std::span<const int> vertices) { return construct(vertices, false); }
	const std::vector<Node> &build_sorted(std::span<const int> vertices) { return construct(vertices, true); }
};

#ifndef AUXILIARY_TREE_HPP
#define AUXILIARY_TREE_HPP

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

template <class Tree> class AuxiliaryTree {
  public:
	struct Node {
		int vertex, parent, length;
		bool selected;
	};

  private:
	struct Frame {
		int vertex, child;
		bool selected;
	};
	const Tree *tree_;
	std::vector<std::pair<int, int>> keys_;
	std::vector<Frame> stack_;
	std::vector<Node> nodes_;
	int finish() {
		Frame current = stack_.back();
		stack_.pop_back();
		int id = size();
		for (int child = current.child; child != -1;) {
			int next = nodes_[child].parent;
			nodes_[child].parent = id;
			nodes_[child].length = tree_->depth[nodes_[child].vertex] - tree_->depth[current.vertex];
			child = next;
		}
		nodes_.push_back({current.vertex, -1, 0, current.selected});
		return id;
	}
	void attach(int child) {
		nodes_[child].parent = stack_.back().child;
		stack_.back().child = child;
	}
	const std::vector<Node> &construct(std::span<const int> vertices, bool sorted) {
		keys_.clear();
		for (int v : vertices) keys_.emplace_back(tree_->subtree(v).first, v);
		if (!sorted) {
			if (keys_.size() < 256) std::sort(keys_.begin(), keys_.end());
			else stable_radix_sort(keys_, [](const auto &x) { return static_cast<unsigned>(x.first); });
		}
		else assert(std::is_sorted(keys_.begin(), keys_.end()));
		nodes_.clear();
		stack_.clear();
		int previous = -1;
		for (auto [position, v] : keys_) {
			(void)position;
			if (v == previous) continue;
			previous = v;
			if (!stack_.empty()) {
				int w = tree_->lca(stack_.back().vertex, v);
				while (tree_->depth[stack_.back().vertex] > tree_->depth[w]) {
					int child = finish();
					if (stack_.empty() || tree_->depth[stack_.back().vertex] < tree_->depth[w])
						stack_.push_back({w, -1, false});
					attach(child);
				}
			}
			stack_.push_back({v, -1, true});
		}
		while (stack_.size() > 1) attach(finish());
		if (!stack_.empty()) finish();
		return nodes_;
	}

  public:
	explicit AuxiliaryTree(const Tree &tree) : tree_(&tree) {}
	AuxiliaryTree(const Tree &&) = delete;
	int size() const { return static_cast<int>(nodes_.size()); }
	bool empty() const { return nodes_.empty(); }
	int root() const { return size() - 1; }
	const Node &operator[](int i) const {
		assert(0 <= i && i < size());
		return nodes_[i];
	}
	const std::vector<Node> &nodes() const { return nodes_; }
	std::vector<std::vector<int>> to_graph() const {
		std::vector<std::vector<int>> graph(size());
		std::vector<int> degree(size());
		for (int i = 0; i < size(); ++i) if (nodes_[i].parent >= 0) {
			++degree[i];
			++degree[nodes_[i].parent];
		}
		for (int i = 0; i < size(); ++i) graph[i].reserve(degree[i]);
		for (int i = 0; i < size(); ++i) if (nodes_[i].parent >= 0) {
			int parent = nodes_[i].parent;
			graph[i].push_back(parent);
			graph[parent].push_back(i);
		}
		return graph;
	}
	const std::vector<Node> &build(std::span<const int> vertices) { return construct(vertices, false); }
	const std::vector<Node> &build_sorted(std::span<const int> vertices) { return construct(vertices, true); }
};

#endif

#line 1 "tree/auxiliary_tree.hpp"



#line 1 "components/radix_sort.hpp"



#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>

template <int Bits, class T, class Key>
void stable_radix_sort(std::vector<T> &values, Key &&key) {
	static_assert(1 <= Bits && Bits <= 16);
	constexpr unsigned buckets = 1U << Bits, mask = buckets - 1;
	using K = std::remove_cvref_t<decltype(key(std::declval<const T &>()))>;
	static_assert(std::is_integral_v<K> && std::is_unsigned_v<K> && sizeof(K) <= sizeof(std::uint64_t));
	if (values.size() < 2) return;
	if (values.size() < 256) {
		std::stable_sort(values.begin(), values.end(), [&](const T &a, const T &b) { return key(a) < key(b); });
		return;
	}
	std::uint64_t first = key(values.front()), previous = first, varying = 0;
	bool sorted = true;
	for (std::size_t i = 1; i < values.size(); ++i) {
		std::uint64_t current = key(values[i]);
		varying |= first ^ current;
		sorted &= previous <= current;
		previous = current;
	}
	if (sorted) return;
	std::vector<T> scratch(values.size());
	using Counts = std::conditional_t<(Bits <= 11), std::array<std::size_t, buckets>, std::vector<std::size_t>>;
	Counts count = [] {
		if constexpr (Bits <= 11) return Counts{};
		else return Counts(buckets);
	}();
	for (int shift = 0; shift < 64; shift += Bits) {
		if (!((varying >> shift) & mask)) continue;
		std::fill(count.begin(), count.end(), 0);
		for (const auto &value : values) ++count[(static_cast<std::uint64_t>(key(value)) >> shift) & mask];
		std::size_t sum = 0;
		for (auto &x : count) {
			auto n = x;
			x = sum, sum += n;
		}
		for (auto &&value : values) scratch[count[(static_cast<std::uint64_t>(key(value)) >> shift) & mask]++] = std::move(value);
		values.swap(scratch);
	}
}

template <class T, class Key>
void stable_radix_sort(std::vector<T> &values, Key key) {
	stable_radix_sort<8>(values, std::move(key));
}

namespace radix_sort_detail {
template <int Bits = 8, class T, class Key>
void sort_by_key(std::vector<T> &values, Key key) {
	using K = std::remove_cvref_t<decltype(key(std::declval<const T &>()))>;
	if constexpr (std::is_integral_v<K> && !std::is_same_v<K, bool> && sizeof(K) <= 8 && std::is_default_constructible_v<T>) {
		using U = std::make_unsigned_t<K>;
		constexpr U sign = std::is_signed_v<K> ? U(1) << (sizeof(K) * 8 - 1) : U(0);
		stable_radix_sort<Bits>(values, [&](const T &value) { return U(U(key(value)) ^ sign); });
	} else std::sort(values.begin(), values.end(), [&](const T &a, const T &b) { return key(a) < key(b); });
}
}


#line 6 "tree/auxiliary_tree.hpp"
#include <cassert>
#include <span>
#line 10 "tree/auxiliary_tree.hpp"

template <class Tree> class AuxiliaryTree {
  public:
	struct Node {
		int vertex, parent, length;
		bool selected;
	};

  private:
	struct Frame {
		int vertex, child;
		bool selected;
	};
	const Tree *tree_;
	std::vector<std::pair<int, int>> keys_;
	std::vector<Frame> stack_;
	std::vector<Node> nodes_;
	int finish() {
		Frame current = stack_.back();
		stack_.pop_back();
		int id = size();
		for (int child = current.child; child != -1;) {
			int next = nodes_[child].parent;
			nodes_[child].parent = id;
			nodes_[child].length = tree_->depth[nodes_[child].vertex] - tree_->depth[current.vertex];
			child = next;
		}
		nodes_.push_back({current.vertex, -1, 0, current.selected});
		return id;
	}
	void attach(int child) {
		nodes_[child].parent = stack_.back().child;
		stack_.back().child = child;
	}
	const std::vector<Node> &construct(std::span<const int> vertices, bool sorted) {
		keys_.clear();
		for (int v : vertices) keys_.emplace_back(tree_->subtree(v).first, v);
		if (!sorted) {
			if (keys_.size() < 256) std::sort(keys_.begin(), keys_.end());
			else stable_radix_sort(keys_, [](const auto &x) { return static_cast<unsigned>(x.first); });
		}
		else assert(std::is_sorted(keys_.begin(), keys_.end()));
		nodes_.clear();
		stack_.clear();
		int previous = -1;
		for (auto [position, v] : keys_) {
			(void)position;
			if (v == previous) continue;
			previous = v;
			if (!stack_.empty()) {
				int w = tree_->lca(stack_.back().vertex, v);
				while (tree_->depth[stack_.back().vertex] > tree_->depth[w]) {
					int child = finish();
					if (stack_.empty() || tree_->depth[stack_.back().vertex] < tree_->depth[w])
						stack_.push_back({w, -1, false});
					attach(child);
				}
			}
			stack_.push_back({v, -1, true});
		}
		while (stack_.size() > 1) attach(finish());
		if (!stack_.empty()) finish();
		return nodes_;
	}

  public:
	explicit AuxiliaryTree(const Tree &tree) : tree_(&tree) {}
	AuxiliaryTree(const Tree &&) = delete;
	int size() const { return static_cast<int>(nodes_.size()); }
	bool empty() const { return nodes_.empty(); }
	int root() const { return size() - 1; }
	const Node &operator[](int i) const {
		assert(0 <= i && i < size());
		return nodes_[i];
	}
	const std::vector<Node> &nodes() const { return nodes_; }
	std::vector<std::vector<int>> to_graph() const {
		std::vector<std::vector<int>> graph(size());
		std::vector<int> degree(size());
		for (int i = 0; i < size(); ++i) if (nodes_[i].parent >= 0) {
			++degree[i];
			++degree[nodes_[i].parent];
		}
		for (int i = 0; i < size(); ++i) graph[i].reserve(degree[i]);
		for (int i = 0; i < size(); ++i) if (nodes_[i].parent >= 0) {
			int parent = nodes_[i].parent;
			graph[i].push_back(parent);
			graph[parent].push_back(i);
		}
		return graph;
	}
	const std::vector<Node> &build(std::span<const int> vertices) { return construct(vertices, false); }
	const std::vector<Node> &build_sorted(std::span<const int> vertices) { return construct(vertices, true); }
};