2020西工大校赛 C王国（LCA）

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=2e5+10;
int depth[N];
int fa[N][21];
int ne[N],e[N],h[N],idx;
int n;
int a[N];
vector<int> num[N];
void add(int a,int b){
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
bool cmp(int a,int b){
    return depth[a]>depth[b];
}
void bfs(){
    memset(depth,0x3f,sizeof depth);
    queue<int> q;
    q.push(1);
    depth[0]=0,depth[1]=1;
    while(q.size()){
        int t=q.front();
        q.pop();
        int i;
        for(i=h[t];i!=-1;i=ne[i]){
            int j=e[i];
            if(depth[j]>depth[t]+1){
                depth[j]=depth[t]+1;
                q.push(j);
                fa[j][0]=t;
                for(int k=1;k<=20;k++){
                    fa[j][k]=fa[fa[j][k-1]][k-1];
                }
            }
        }
    }
}
int lca(int a,int b){
    if(depth[a]<depth[b])
        swap(a,b);
    int i;
    for(i=20;i>=0;i--){
        if(depth[fa[a][i]]>=depth[b]){
            a=fa[a][i];
        }
    }
    if(a==b){
        return a;
    }
    for(i=20;i>=0;i--){
        if(fa[a][i]!=fa[b][i]){
            a=fa[a][i];
            b=fa[b][i];
        }
    }
    return fa[a][0];
}
int main(){
    cin>>n;
    int i;
    memset(h,-1,sizeof h);
    for(i=1;i<=n;i++){
        scanf("%d",&a[i]);
        num[a[i]].push_back(i);
    }
    for(i=1;i<n;i++){
        int a,b;
        scanf("%d%d",&a,&b);
        add(a,b);
        add(b,a);
    }
    bfs();
    int res=0;
    for(i=1;i<=n;i++){
        if(num[i].size()){
            sort(num[i].begin(),num[i].end(),cmp);
            int pos=num[i][0];
            for(int j=1;j<num[i].size();j++){
                int p=lca(pos,num[i][j]);
                int k=depth[pos]+depth[num[i][j]]-2*depth[p];
                res=max(res,k*k);
            }
        }
    }
    cout<<res<<endl;
}