けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

木の直径 (AOJ Course GRL_5_A)

めちゃめちゃ簡単な実装で良いことがわかったので。
なんか、DFS 2 回か、全方位やるかなのかと思ってたけど、とても簡単な木DPで直径が求められることを知った。

問題へのリンク

問題概要

 N 頂点の重み付き木が与えられるので、その直径の長さを求めよ。

制約

  •  2 \le N \le 10^{5}

解法

全方位木 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;
    }
}