IOI 13 Dreaming ( 樹形DP )

PEG Judge - IOI '13 - Dreaming
題目大意是給一個森林，希望用任意數量的長度為L的邊將他黏成一棵樹，使得最後這棵樹的直徑最小，求的是這樣的直徑。
很容易想到，最理想的方法一定是將所有樹折半，以他的樹心( 該點對於樹上其他點的最長距離最小 ) 連結，而且是將所有的樹的樹心連到半徑( 樹心到該樹上最遠點的距離 ) 最大的樹的樹心上。這麼一來，合成後的樹的直徑的最大值一定是 max{ 最大樹的直徑，最大樹的半徑 + L + 次大樹的半徑，第三大樹半徑 + L + L + 次大樹半徑 }。
所以這個題目的重點是要如何求出三大樹的半徑。我們可以考慮用樹形DP，每一棵樹分別考慮。一個重要的啟發是，如果一棵樹是有根的，那麼任何一個節點的最遠距離要嘛是他父親的往上的最遠距離加上轉移邊權，或者是他父親往下( 往當前節點的兄弟走 )的最遠距離加上轉移邊權。所以我們首先提起任何一個節點使其成為有根樹，先做一次DFS維護兩樣東西：由這個點開始只往下走到底的最長距離，次長距離和走最長距離要選的兒子。注意實作中如果有多個最長距離，自然而然次長距離會等於最長距離所以不會出問題。接著做第二次DFS由上往下依次更新往上的最遠距離和該點的最遠距離。
我實作中沒弄清楚的細節有三個，一個是忘了幫根節點計算最遠點，另一個是因為半徑初始化用了INF，但有可能某個樹是只有一個節點的，這時候沒更新到就慘了，最後一個是忘了判少於三棵樹的情況。

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1e5 + 5;
const int MAXM = MAXN;
const int MAXL = 1e4 + 4;
const int INF = 0x3f3f3f3f;

int n, m, l;

struct Edge{
    int u, v, w;
    Edge(int _u = 0, int _v = 0, int _w = 0): u(_u), v(_v), w(_w){}
} es[ MAXM * 2 ];

vector<int> G[ MAXN ];

int cc_cnt, cc[ MAXN ];
int bfur_dis[ MAXN ], bsfur_dis[ MAXN ]; // furthest distance to bottom, second furthest to bottom
int bfdc[ MAXN ]; // child to furthest distance to bottom ( leaf )
void dfs1(int u, int fa){ // find furthest and second furthest distance of node towards bottom
    cc[ u ] = cc_cnt;
    for(int vidx: G[ u ]){
        Edge &e = es[ vidx ];
        if( e.v == fa ) continue;
        dfs1( e.v, u );
        if( bfur_dis[ u ] < e.w + bfur_dis[ e.v ] ){
            bfur_dis[ u ] = e.w + bfur_dis[ e.v ];
            bfdc[ u ] = e.v;
        }
    }
    for(int vidx: G[ u ]){
        Edge &e = es[ vidx ];
        if( e.v == fa ) continue;
        if( e.v == bfdc[ u ] ) continue;
        if( bsfur_dis[ u ] < e.w + bfur_dis[ e.v ] )
            bsfur_dis[ u ] = e.w + bfur_dis[ e.v ];
    }
}

int ufur_dis[ MAXN ]; // furthest distance towards top ( root )
void dfs2(int u, int fa, int &diameter, int &radius){
    for(int vidx: G[ u ]){
        Edge &e = es[ vidx ];
        if( e.v == fa ) continue;

        if( bfdc[ u ] == e.v )
            ufur_dis[ e.v ] = max( ufur_dis[ u ], bsfur_dis[ u ] ) + e.w;
        else
            ufur_dis[ e.v ] = max( ufur_dis[ u ], bfur_dis[ u ]  ) + e.w;

        int v_fur_dis = max( ufur_dis[ e.v ], bfur_dis[ e.v ] );
        diameter = max( diameter, v_fur_dis );
        radius = min( radius, v_fur_dis );
        dfs2( e.v, u, diameter, radius );
    }
}

void solve(){
    int max_diameter = 0;
    vector<int> rad;
    for(int u = 0; u < n; ++u) if( !cc[ u ] ){
        ++cc_cnt;
        dfs1( u, -1 );
        // int diameter = 0, radius = INF; !!OOPS!!
        int diameter = bfur_dis[ u ], radius = bfur_dis[ u ];
        dfs2( u, -1, diameter, radius );
        max_diameter = max( diameter, max_diameter );
        if( radius < INF ) rad.push_back( radius ); // if radius == INF, node is alone
    }
    sort( rad.begin(), rad.end(), greater<int>() );
    int ans = max_diameter;
    if( rad.size() >= 3 ) // careful, maybe there aren't enough nodes
        ans = max( ans, rad[ 1 ] + l + l + rad[ 2 ] );
    if( rad.size() >= 2 )
        ans = max( ans, rad[ 0 ] + l + rad[ 1 ] );
    printf("%d\n", ans);
}

int main(){
    scanf("%d%d%d", &n, &m, &l);
    for(int i = 0; i < m; ++i){
        int u, v, w; scanf("%d%d%d", &u, &v, &w);
        es[ i ] = Edge( u, v, w );
        es[ m + i ] = Edge( v, u, w );
        G[ u ].push_back( i );
        G[ v ].push_back( m + i );
    }
    solve();
    return 0;
}