给出\(n\le 2\cdot 10^5\)的一棵树,每个节点有一个颜色。
求出路径长度为\(2\)路径首端和尾端拥有相同颜色,且路径上其他点不存在相同颜色的点的路径条数。
当时看错题了,把颜色抽出来后没法做了。
后来感觉能点分治,然后把题看对了,遂写了一个极其抽象的点分治。
除此之外,把某种颜色抽出来后,建立出来虚树,在树上进行简单的\(dp\)即可。
由于虚树大小和每次拿出来的点数同阶,故复杂度是线性的。
除此之外,考虑另一种写法:不用建立虚树直接在树上按照虚树的思想进行dp。
每次只考虑子树内的颜色答案统计。其他颜色不用管,乘上对应组合数即可。
const int MAXN=300010;
int T;
int n,m,len;
int a[MAXN];
ll ans;
int lin[MAXN],ver[MAXN<<1],nex[MAXN<<1];
int c[MAXN];
inline void add(int x,int y)
{
ver[++len]=y;
nex[len]=lin[x];
lin[x]=len;
}
void dp(int x,int fa)
{
++c[a[x]];
int las=c[a[x]];
go(x)
{
if(tn==fa)continue;
dp(tn,x);
int w=c[a[x]]-las;
ans+=w+(ll)w*(w-1)/2;
c[a[x]]=las;
}
}
int main()
{
freopen("1.in","r",stdin);
sc(T);
while(T--)
{
len=ans=0;
rep(1,n,i)lin[i]=0,c[i]=0;
sc(n);
rep(1,n,i)sc(a[i]);
rep(2,n,i)
{
int x,y;
sc(x);sc(y);
add(x,y);add(y,x);
}
dp(1,0);
rep(1,n,i)ans+=(ll)c[i]*(c[i]-1)/2;
putl(ans);
}
return 0;
}
标签:Educational,颜色,int,分治,Codeforces,Round,MAXN,虚树,dp
From: https://www.cnblogs.com/chdy/p/18052302