BZOJ2870 最长道路tree(并查集+LCA)

题意

(n<=50000)

题解

#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
#include<algorithm>
using namespace std;
const long long N=60000;
long long cnt,head[N];
long long sum[N],dep[N],f[N][25];
long long fa[N];
long long n,a[N],b[N],book[N],ans;
struct edge{
    long long to,nxt;
}e[N*2];
struct node{
    long long w,id;
}c[N];
void add(long long u,long long v){
    cnt++;
    e[cnt].nxt=head[u];
    e[cnt].to=v;
    head[u]=cnt;
}
void dfs(long long u,long long ff,long long deep,long long w){
    sum[u]=w;
    f[u][0]=ff;
    dep[u]=deep;
    for(long long i=head[u];i;i=e[i].nxt){
        long long v=e[i].to;
        if(v==ff)continue;
        dfs(v,u,deep+1,w+1);
    }
}
bool cmp(node a,node b){
    return a.w>b.w;
}
long long find(long long x){
    if(fa[x]==x)return x;
    else return fa[x]=find(fa[x]);
}
long long getlca(long long x,long long y){
    if(dep[x]<dep[y])swap(x,y);
    for(long long i=20;i>=0;i--){
        if(dep[f[x][i]]>=dep[y]){
            x=f[x][i];
        }
    } 
    if(x==y)return x;
    for(long long i=20;i>=0;i--){
        if(f[x][i]!=f[y][i]){
            x=f[x][i];y=f[y][i];
        }
    }
    return f[x][0];
}
long long getline(long long x,long long y){
    long long LCA=getlca(x,y);
    return sum[x]+sum[y]-sum[LCA]*2+1; 
}
void merge(long long x,long long y){
    long long tmp=0,cx,cy,len;
    len=getline(a[x],b[x]);if(len>tmp){cx=a[x];cy=b[x];tmp=len;}
    len=getline(a[y],b[y]);if(len>tmp){cx=a[y];cy=b[y];tmp=len;}
    len=getline(a[x],a[y]);if(len>tmp){cx=a[x];cy=a[y];tmp=len;}
    len=getline(b[x],b[y]);if(len>tmp){cx=b[x];cy=b[y];tmp=len;}
    len=getline(a[x],b[y]);if(len>tmp){cx=a[x];cy=b[y];tmp=len;}
    len=getline(a[y],b[x]);if(len>tmp){cx=a[y];cy=b[x];tmp=len;}
    fa[y]=x;
    a[x]=cx;
    b[x]=cy;
}
int main(){
    scanf("%lld",&n);
    for(long long i=1;i<=n;i++){
        scanf("%lld",&c[i].w);
        c[i].id=i;
    }
    for(long long i=1;i<n;i++){
        long long u;long long v;
        scanf("%lld%lld",&u,&v);
        add(u,v);
        add(v,u);
    }
    for(long long i=1;i<=n;i++){
        a[i]=b[i]=i;fa[i]=i;
    }
    dfs(1,0,1,0);
    for(long long i=1;i<=20;i++){
        for(long long j=1;j<=n;j++){
            f[j][i]=f[f[j][i-1]][i-1];
        }
    }
    sort(c+1,c+1+n,cmp);
    for(long long i=1;i<=n;i++){
        long long u=c[i].id;
        book[u]=1;
        for(long long j=head[u];j;j=e[j].nxt){
            long long v=e[j].to;
            if(book[v]==0)continue;
            merge(find(u),find(v));
        }
        ans=max(ans,c[i].w*getline(a[u],b[u])); 
    }
    printf("%lld",ans);
    return 0;
}