牛客OI周賽15-提高組 A 環球旅行-樹的直徑

題目鏈接：https://ac.nowcoder.com/acm/contest/4912/A

#include <bits/stdc++.h>
#define LL long long
using namespace std;
 
struct E{
    int to, w;
    E(int a, int b){
        to=a; w=b;
    }
};
vector<E> e[1000005], Edge[1000005];
int k, d, pos1, pos2;
int vis[1000005];
LL f[1000005], fmx[2][1000005];
 
void dfs(int u, int fa, int s){
    if(s>=d){
        d=s, k=u;
    }
    for(int i=0;i<e[u].size();i++){
        int v=e[u][i].to, w=e[u][i].w;
        if(v!=fa){
            dfs(v, u, s+w);
        }
    }
}
 
int DFS(int u, int fa, int root){//dfs出直徑鏈建新圖
    if(u==root){
        return 1;
    }
    for(auto x: e[u]){
        if(x.to!=fa&&DFS(x.to, u, root)){
            Edge[u].push_back(E{x.to, x.w});
            Edge[x.to].push_back(E{u, x.w});
            vis[x.to]=1; return 1;
        }
    }
    return 0;
}
 
LL DP(int u){//子樹的最長鏈
 
    vis[u]=1;
    for(auto x: e[u]){
        if(vis[x.to]==0){
            f[u]=max(f[u], DP(x.to)+x.w);
        }
    }
    return f[u];
}
 
void Dfs(int u, int fa, int pos, int s){//在直徑上DP出所有點的最大鏈長,並且求前綴max
 
    DP(u);
    fmx[pos][u]=max(fmx[pos][fa], s+f[u]);
    for(auto x: Edge[u]){
        if(x.to!=fa){
            //cout<<u<<" "<<fmx[pos][u]<<endl;
            Dfs(x.to, u, pos, s+x.w);
        }
    }
}
 
int ans=1e9+7;
void getans(int u, int fa){
    for(auto x: Edge[u]){
        if(x.to!=fa){
            ans=min(1ll*ans, max(fmx[0][u], fmx[1][x.to]));
            getans(x.to, u);
        }
    }
}
 
int main()
{
    int n, u, v, w;
    scanf("%d", &n);
    for(int i=1; i<n; i++){
        scanf("%d%d%d",&u,&v,&w);
        e[u].push_back(E(v, w));
        e[v].push_back(E(u, w));
    }
    pos1=pos2=k=0;
    dfs(1, 0, 0); pos1=k;
    dfs(k, 0, 0); pos2=k;
    vis[pos1]=1;
    DFS(pos1, 0, pos2);
    Dfs(pos1, 0, 0, 0);
    Dfs(pos2, 0, 1, 0);
    getans(pos1, 0);
    printf("%d\n", ans);
 
    return 0;
}