AcWing 1072. 树的最长路径 --树的直径模板

点击跳转例题

思路：

我们知道：树上 任意两点 的路径是 唯一 确定的，因此我们可以暴力枚举 起点 和 终点 找出最长路径

如果这样做的话，我们来思考一下时间复杂度：

枚举 起点 和 终点 — O(n^2)
找出两点之间的路径长度 — O(logn)

我们考虑换一种 枚举方式：枚举路径的 起点和终点 → 枚举路径的 中间节点
每次记录以当前节点为树的根下的最长链和次长链即可。

 

#include <bits/stdc++.h>
#define int long long //(有超时风险)
#define PII pair<int,int>
#define endl '\n'
#define all(x) x.begin(), x.end()

using namespace std;

const int N=2e6+10,M=1e3+10,mod=998244353,INF=0x3f3f3f3f;

int a[N];
int h[N],e[N],ne[N],w[N],idx;

void add(int a,int b,int c)
{
    w[idx]=c;e[idx]=b;ne[idx]=h[a];h[a]=idx++;
}

int ans=0;
int dfs(int x,int fa)
{
    int dist=0;//返回的是以x节点的最长链的长度；
    //d1,d2也可以用数组记录。
    int d1=0,d2=0;//表示这个点的最长链和次长链；

    for(int i=h[x];~i;i=ne[i])
    {
        int j=e[i];
        if(j==fa)continue;
        //加上w[i],是因为，e[i]存储的为x->e[i]这个点的编号，w[i]存储的是x->e[i]这条边的权值；
        int d=dfs(j,x)+w[i];//x的其中一个子节点的最长链，加上本身这条边的权值
        
        if(d>d1)d2=d1,d1=d;
        else if(d>d2)d2=d;
        
        dist=max(d,dist);//所有的子链取max
    }
    //所有的：挂在x为节点的，最长链
    ans=max(ans,d1+d2);
    return dist;
}

//划分集合：所有以某个节点为挂点的最长链的方案，属性为最大值；
signed main()
{
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n;cin>>n;
    memset(h,-1,sizeof h);idx=0;
    for(int i=1;i<=n-1;i++)
    {
        int x,y,z;cin>>x>>y>>z;
        add(x,y,z);
        add(y,x,z);
    }
    dfs(1,-1);
    cout<<ans<<endl;
    return 0;
}