めちゃめちゃ簡単な実装で良いことがわかったので。
なんか、DFS 2 回か、全方位やるかなのかと思ってたけど、とても簡単な木DPで直径が求められることを知った。
問題概要
頂点の重み付き木が与えられるので、その直径の長さを求めよ。
制約
解法
全方位木 DP するときは、各頂点を根としたときの、各部分木に対して「その部分木の根から葉への距離の最大値」を求める感じ。
なんだけど、単純に根をどこか適当に決めて、普通の木DPでできるっぽい
- dp[ v ] := v を根とする部分木において、v から葉までの距離の最大値
としてあげればこれは普通の木DPで求められる。そしてメモ化再帰内で、
- 答えを表す変数 and を、頂点 v の各子供についての dp の値を大きい順にソートして上位 2 個をとった合計値と比べて更新する
という風にする感じ。
#include <bits/stdc++.h> using namespace std; template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; } template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; } using Edge = pair<int, long long>; using Graph = vector<vector<Edge>>; int N; Graph G; long long rec(int v, int p, long long &ans) { vector<long long> chs; for (auto e : G[v]) { if (e.first == p) continue; chs.push_back(rec(e.first, v, ans) + e.second); } sort(chs.begin(), chs.end(), greater<long long>()); if (chs.size() == 0) return 0; else if (chs.size() == 1) chmax(ans, chs[0]); else chmax(ans, chs[0] + chs[1]); return chs[0]; } int main () { while (cin >> N) { G.assign(N, vector<Edge>()); for (int i = 0; i < N-1; ++i) { int u, v; long long w; cin >> u >> v >> w; G[u].push_back({v, w}); G[v].push_back({u, w}); } long long ans = 0; rec(0, -1, ans); cout << ans << endl; } }