考虑到树上操作;
首先题目要我们求每条路径上出现不同颜色的数量,并把所有加起来得到答案;
我们知道俩俩点之间会形成一条路径,所以我们可以知道每个样例的总的路径的数目为:n*(n-1)/2;
这样单单的求,每条路径(n:2e5)无疑会爆;
这样我们假设所有路径上都存在所有的颜色,所有总的答案为n*(n-1)/2*n;
然后我们再在里面减去我们不需要的;
这里我们要运用虚树(当前图的信息整合而已)的思想,其实也没有建出一颗树;
对于一个顶点u,颜色为x,在它的子树内所有以颜色x为根的子树都要舍去;//这个过程用dfs实现
代码又注释其中过程
#include<iostream>
#include<vector>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<vector>
using namespace std;
typedef long long ll;
const int M=2e5+5;
vector < int > e[M];
int top[M];此时颜色为i的联通块的top
int col[M];
int a[M];//以i为根节点的子树中,不在i所在的联通块中节点的number
int b[M];//以颜色i为分界线,不在根节点所在的联通块中的节点个数
ll n,ans;
int dfs(int u,int f){
int u_sz=1,curtop=top[col[u]];
top[col[u]]=u;
for(int i=0;i<e[u].size();i++){
int v=e[u][i];
if(v!=f){
a[u]=0;
int v_sz=dfs(v,u);
ll x=v_sz-a[u];
ans-=x*(x-1)/2;
u_sz+=v_sz;
}
}
if(curtop){
a[curtop]+=u_sz;
}
else
b[col[u]]+=u_sz;
top[col[u]]=curtop;
return u_sz;
}
int main(){
int sign=1;
while(~scanf("%lld",&n)){
for(int i=0;i<=n;i++){
e[i].clear();
a[i]=0,b[i]=0;
top[i]=0;
}
for(int i=1;i<=n;i++)
scanf("%d",&col[i]);
for(int i=1;i<n;i++){
int x,y;
scanf("%d%d",&x,&y);
e[x].push_back(y);
e[y].push_back(x);
}
ans=n*(n-1)/2*n;
// cout<<ans<<endl;
dfs(1,-1);
for(int i=1;i<=n;i++){//这一步是因为dfs中没有对根节点进行操作
ll x=n-b[i];
ans-=x*(x-1)/2;
}
printf("Case #%d: %lld\n",sign++,ans);
}
return 0;
}