Educational Codeforces Round 162 E 点分治 虚树 树形dp

传送门

给出\(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;
}