test/Diameter.test.cpp

View this file on GitHub
Last update: 2021-03-22 17:28:05+09:00
Problem: https://judge.yosupo.jp/problem/tree_diameter

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
#include "./../Graph/Diameter.cpp"
#include "./../Graph/DiameterPath.cpp"
#include <iostream>
using namespace std;

int main() {
	cin.tie(nullptr);
	ios::sync_with_stdio(false);
	int n;
	cin >> n;
	Graph g(n);
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		Weight c;
		cin >> a >> b >> c;
		g[a].emplace_back(b, c);
		g[b].emplace_back(a, c);
	}

	auto [dist, u, v] = Diameter(g);
	auto [_dist, path] = DiameterPath(g);
	cout << dist << ' ' << path.size() << '\n';
	for (size_t i = 0; i < path.size(); i++) {
		cout << path[i] << (i != path.size() - 1 ? ' ' : '\n');
	}
}

#line 1 "test/Diameter.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
#line 2 "Graph/GraphTemplate.cpp"
#include <vector>
#include <utility>
#include <iostream>
#include <limits>

using Weight = long long;
constexpr Weight INF = std::numeric_limits<Weight>::max();
struct Edge {
	int to;
	Weight cost;
	Edge() : to(-1), cost(-1) {}
	Edge(int _to, Weight _cost = 1) : to(_to), cost(_cost) {}
	friend bool operator<(const Edge& e1, const Edge& e2) {
		return e1.cost < e2.cost;
	}
	friend bool operator>(const Edge& e1, const Edge& e2) {
		return e1.cost > e2.cost;
	}
	friend std::ostream& operator<<(std::ostream& os, const Edge& e) {
		return os << "->" << e.to << '(' << e.cost << ')';
	}
};
using UnWeightedGraph = std::vector<std::vector<int>>;
using Graph = std::vector<std::vector<Edge>>;
struct Edge2 {
	int from, to;
	Weight cost;
	Edge2() : from(-1), to(-1), cost(0) {}
	Edge2(int _from, int _to, Weight _cost) : from(_from), to(_to), cost(_cost) {}
	friend bool operator<(const Edge2& e1, const Edge2& e2) {
		return e1.cost < e2.cost;
	}
	friend bool operator>(const Edge2& e1, const Edge2& e2) {
		return e1.cost > e2.cost;
	}
	friend std::ostream& operator<<(std::ostream& os, const Edge2& e) {
		return os << e.from << "->" << e.to << '(' << e.cost << ')';
	}
};
using UnWeightedEdges = std::vector<std::pair<int, int>>;
using Edges = std::vector<Edge2>;
using Matrix = std::vector<std::vector<Weight>>;

auto add_edge(UnWeightedGraph& graph, int v, int u) {
	graph[v].push_back(u);
	graph[u].push_back(v);
}
auto add_edge(Graph& graph, int v, int u, Weight cost) {
	graph[v].emplace_back(u, cost);
	graph[u].emplace_back(v, cost);
}
auto to_graph(const UnWeightedGraph& graph, Weight cost = 1) {
	Graph result(graph.size());
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (int v : graph[i]) {
			result[i].emplace_back(v, cost);
		}
	}
	return result;
}
auto to_unweighted_graph(const Graph& graph) {
	UnWeightedGraph result(graph.size());
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (auto [v, cost] : graph[i]) {
			result[i].push_back(v);
		}
	}
	return result;
}
auto to_edges(const UnWeightedGraph& graph, bool unique = false) {
	std::vector<std::pair<int, int>> edges;
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (int v : graph[i]) {
			if (!unique || static_cast<int>(i) < v) edges.emplace_back(i, v);
		}
	}
	return edges;
}
auto to_edges(const Graph& graph) {
	Edges edges;
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (auto [v, cost] : graph[i]) {
			edges.emplace_back(i, v, cost);
		}
	}
	return edges;
}
#line 4 "Graph/Diameter.cpp"
#include <tuple>

std::tuple<Weight, int, int> Diameter(const Graph& graph) {
	auto dfs = [&](auto&& f, int v, int p, Weight depth) -> std::pair<Weight, int> {
		std::pair<Weight, int> result(depth, v);
		for (auto e : graph[v]) {
			if (e.to != p) {
				auto tmp = f(f, e.to, v, depth + e.cost);
				if (tmp.first > result.first) result = tmp;
			}
		}
		return result;
	};
	auto d1 = dfs(dfs, 0, -1, 0);
	auto d2 = dfs(dfs, d1.second, -1, 0);
	return {d2.first, d1.second, d2.second};
}
Weight DiameterLength(const Graph& graph) {
	return std::get<0>(Diameter(graph));
}
#line 4 "Graph/DiameterPath.cpp"
#include <algorithm>
#line 7 "Graph/DiameterPath.cpp"
#include <functional>

std::tuple<Weight, std::vector<int>> DiameterPath(const Graph& graph) {
	int n = graph.size();
	std::vector<Weight> dist0(n);
	auto dfs = [&](auto self, int v, int p, Weight d) -> void {
		dist0[v] = d;
		for (const auto& u : graph[v])
			if (u.to != p) {
				self(self, u.to, v, d + u.cost);
			}
	};
	dfs(dfs, 0, -1, 0);

	int s = std::max_element(dist0.begin(), dist0.end()) - dist0.begin();
	std::vector<Weight> dist(n);
	std::vector<int> par(n);
	auto dfs2 = [&](auto self, int v, int p, Weight d) -> void {
		dist[v] = d;
		par[v] = p;
		for (const auto& u : graph[v]) {
			if (u.to != p) {
				self(self, u.to, v, d + u.cost);
			}
		}
	};
	dfs2(dfs2, s, -1, 0);
	auto t = std::max_element(dist.begin(), dist.end());
	std::vector<int> path{static_cast<int>(t - dist.begin())};
	for (int p = 0; (p = par[path.back()]) != -1;) {
		path.push_back(p);
	}
	std::reverse(path.begin(), path.end());
	return {*t, path};
}
#line 5 "test/Diameter.test.cpp"
using namespace std;

int main() {
	cin.tie(nullptr);
	ios::sync_with_stdio(false);
	int n;
	cin >> n;
	Graph g(n);
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		Weight c;
		cin >> a >> b >> c;
		g[a].emplace_back(b, c);
		g[b].emplace_back(a, c);
	}

	auto [dist, u, v] = Diameter(g);
	auto [_dist, path] = DiameterPath(g);
	cout << dist << ' ' << path.size() << '\n';
	for (size_t i = 0; i < path.size(); i++) {
		cout << path[i] << (i != path.size() - 1 ? ' ' : '\n');
	}
}