对两个错两个T一个，但算法是线性的，也用了快读快输。基础版的是60分，后面大数据反而过了1，3两点错了，该开longlong的也开了，实在找不到哪里有问题。

#include<bits/stdc++.h>
using namespace std;
int n,a,b,cnt=0,jl[3000010]={0},col[3000010]={0},sum[3000010]={0},head[6000010]={0};
int js[3000010]={0},bian[6000010]={0};
bool mark[3000010]={0};
struct node{int to,next;}tu[6000010]={0};
vector<int> v[3000010];
inline int r(){
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9')ch=getchar();
    while(ch>='0'&&ch<='9')x=(x<<1)+(x<<3)+(ch^48),ch=getchar();
    return x*f;
}
inline void w(long long X){
	int s[20],top=0;
	while(X) {s[++top]=X%10; X/=10;}
	if(!top) s[++top]=0;
	while(top) putchar(s[top--]+'0');
}
int pre_dfs(int x,int fa){sum[col[x]]++;
    int num=1,temp,pre_jl=sum[col[x]],last=0;    
    for(int i=head[x];i;i=tu[i].next){
        if(tu[i].to==fa)continue;
        temp=pre_dfs(tu[i].to,x);
        if(pre_jl!=sum[col[x]])last=tu[i].to;
        bian[i]=temp;num+=temp;
        if(i%2)bian[i+1]=n-temp;
        else bian[i-1]=n-temp;//bian是记录该边所连的点为根的子树大小，去的bian的大小就是子树大小，回来的就是n-子树大小。边从1开始存，所以要这样处理。
    }
    if(mark[col[x]])    return num;
    else if(sum[col[x]]==jl[col[x]]&&pre_jl==1)v[col[x]].push_back(x),js[x]=last;//如果本来一个点都没扫到，现在扫到的点数又等于所有该颜色点的个数，那这个点一定是一个端点
    else if(sum[col[x]]==pre_jl)v[col[x]].push_back(x),js[x]=fa;//如果下面没有该颜色的点，那这个点也是端点
    //js纪录的是通往链中间那坨的点，在后面计算答案要排除的
    if(v[col[x]].size()>2)mark[col[x]]=1;
    return num;
}
void solve_1(int x){
    queue<int> q;long long sum=0,ans=0,temp;
    for(int i=head[x];i;i=tu[i].next)
        q.push(bian[i]),sum+=bian[i];
    while(!q.empty()){
        temp=q.front();q.pop();
        sum-=temp;ans+=sum*temp;
    }w(ans+(n-1));cout<<endl;
}//所有子树两两乘积和+n-1
void solve_2(int x,int y){
    long long sum_x=1,sum_y=1;
    for(int i=head[x];i;i=tu[i].next)
        if(tu[i].to!=js[x])sum_x+=bian[i];
    for(int i=head[y];i;i=tu[i].next)
        if(tu[i].to!=js[y])sum_y+=bian[i];
    w(sum_x*sum_y);cout<<endl;
}//去掉中间一坨点后，剩下两颗子树的大小乘积
void add(int u,int v){
    cnt++;tu[cnt].to=v;tu[cnt].next=head[u];head[u]=cnt;
    cnt++;tu[cnt].to=u;tu[cnt].next=head[v];head[v]=cnt;
}
int main(){
    n=r();for(int i=1;i<=n;i++)col[i]=r(),jl[col[i]]++;//记录每种颜色的点数
    for(int i=1;i<n;i++)a=r(),b=r(),add(a,b);
    pre_dfs(1,0);for(int i=1;i<=n;i++){
        if(mark[i])cout<<0<<endl;
        if(v[i].size()==1)solve_1(v[i][0]);
        if(v[i].size()==0){w((1ll*n*(n-1))/2);cout<<endl;}
        if(v[i].size()==2)solve_2(v[i][0],v[i][1]);
        //分类讨论
    }
    system("pause");return 0;
}