木の直径

tree/tree_diameter.hpp · GitHub で見る

共通ページ

重み付き無向木で、重みの合計が最大になる単純パスと、その長さを求めます。負の重みも扱えます。

#include "lib/tree/tree_diameter.hpp"

TreeDiameter<> tree(4);
tree.add_edge(0, 1, 3);
tree.add_edge(1, 2, 4);
tree.add_edge(1, 3, 2);
auto result = tree.solve();
assert(result.distance == 7);
// result.path は 0, 1, 2 またはその逆順

操作・メンバー	内容・返り値	計算量
TreeDiameter<T>(n = 0)	頂点数を指定。T の既定値は int64_t	\(O(n)\)
add_edge(u,v,weight = 1)	無向辺を追加し、0 始まりの辺 ID を返す	償却 \(O(1)\)
solve()	Result{distance,path} を返す	\(O(n)\)
result.distance	直径の重みの合計（T）	\(O(1)\)
result.path	両端を含むパスの頂点列（vector<int>）	\(O(1)\)
size() / edge_count()	頂点数／辺数	\(O(1)\)
edges()	追加順の Edge{from,to,weight} の配列への const 参照	\(O(1)\)
reserve(m)	辺 m 本分の領域を予約	\(O(m)\)
clear()	全辺を削除。頂点数は維持	\(O(n)\)

頂点番号は [0,n)、頂点数は INT_MAX/2 以下です。solve() の前に閉路のない連結な木を作り、非空なら辺数を n-1 にしてください。reserve(m) は 0 <= m <= n-1 の範囲で指定します。空の木では m=0 のみです。

単頂点のパスを長さ 0 として含めます。単点の木や全ての辺が負の木では、長さ 0 の単頂点パスを返します。空の木では長さ 0 と空列を返します。最適なパスが複数ある場合の選択・向きは保証しません。

T には重みと全てのパスの長さを表現できる型を指定します。128 bit 整数が必要なら TreeDiameter<__int128_t> を使ってください。solve() は入力を変更せず、繰り返し呼べます。コピーと返された結果は独立します。

頂点数を \(n\) として、構築・直径の計算は \(O(n)\) 時間、入力・結果・作業領域を含めて \(O(n)\) 空間です。

検証コード

• test/tree_diameter.test.cpp

コード

#ifndef TREE_DIAMETER_HPP
#define TREE_DIAMETER_HPP

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

template<class T = std::int64_t> class TreeDiameter {
	int n_;

  public:
	struct Edge { int from, to; T weight; };
	struct Result { T distance{}; std::vector<int> path; };

  private:
	std::vector<Edge> edges_;

  public:
	explicit TreeDiameter(int n = 0) : n_(n) {
		assert(0 <= n && n <= std::numeric_limits<int>::max() / 2);
		edges_.reserve(n ? n - 1 : 0);
	}
	TreeDiameter(const TreeDiameter &) = default;
	TreeDiameter &operator=(const TreeDiameter &other) {
		if (this != &other) { TreeDiameter copy(other); swap(copy); }
		return *this;
	}
	TreeDiameter(TreeDiameter &&other) noexcept : n_(std::exchange(other.n_, 0)), edges_(std::move(other.edges_)) {}
	TreeDiameter &operator=(TreeDiameter &&other) noexcept {
		if (this != &other) n_ = std::exchange(other.n_, 0), edges_ = std::move(other.edges_);
		return *this;
	}
	void swap(TreeDiameter &other) noexcept {
		std::swap(n_, other.n_);
		edges_.swap(other.edges_);
	}
	int size() const { return n_; }
	int edge_count() const { return static_cast<int>(edges_.size()); }
	void reserve(std::size_t m) {
		assert(m <= std::size_t(n_ ? n_ - 1 : 0));
		edges_.reserve(m);
	}
	int add_edge(int u, int v, T weight = T(1)) {
		assert(0 <= u && u < n_ && 0 <= v && v < n_ && u != v);
		assert(edges_.size() < std::size_t(n_ - 1));
		int id = edge_count();
		edges_.push_back({u, v, std::move(weight)});
		return id;
	}
	const std::vector<Edge> &edges() const { return edges_; }
	void clear() { edges_.clear(); }
	[[nodiscard]] Result solve() const {
		assert(edges_.size() == std::size_t(n_ ? n_ - 1 : 0));
		Result result;
		if (!n_) return result;
		struct State { int degree = 0, link = 0, parent = -1, endpoint = 0; T best{}; };
		std::vector<State> state(n_);
		for (int id = 0; id < edge_count(); ++id) {
			const auto &e = edges_[id];
			++state[e.from].degree, ++state[e.to].degree;
			state[e.from].link ^= id, state[e.to].link ^= id;
		}
		std::vector<int> order;
		order.reserve(n_ - 1);
		for (int v = 0; v < n_; ++v) {
			state[v].endpoint = v;
			if (v && state[v].degree == 1) order.push_back(v);
		}
		int from = 0, to = 0, middle = 0;
		for (std::size_t i = 0; i < order.size(); ++i) {
			int v = order[i], id = state[v].link;
			const auto &e = edges_[id];
			int p = e.from ^ e.to ^ v;
			state[v].parent = p;
			T candidate = state[v].best + e.weight;
			T joined = state[p].best + candidate;
			if (result.distance < joined) {
				result.distance = std::move(joined);
				from = state[v].endpoint, to = state[p].endpoint, middle = p;
			}
			if (state[p].best < candidate) state[p].best = std::move(candidate), state[p].endpoint = state[v].endpoint;
			state[p].link ^= id;
			if (--state[p].degree == 1 && p) order.push_back(p);
		}
		assert(order.size() == std::size_t(n_ - 1));
		for (int v = from; v != middle; v = state[v].parent) result.path.push_back(v);
		result.path.push_back(middle);
		std::size_t first = result.path.size();
		for (int v = to; v != middle; v = state[v].parent) result.path.push_back(v);
		std::reverse(result.path.begin() + first, result.path.end());
		return result;
	}
};

#endif

Back to top page