CF1923E Count Paths 题解

CF1923E Count Paths

点分治模板题。

假设当前处理的树根为 $x$，我们考虑如何统计经过点 $x$ 的合法路径。

$1$：存在一个与 $x$ 颜色相同的点，且这个点到 $x$ 的路径上没有与之颜色相同的点，那么这个点和 $x$ 就构成了一条合法路径。

$2$：存在一个与 $x$ 颜色不相同的点，且这个点到 $x$ 的路径上没有与之颜色相同的点。那么这个点与 $x$ 的其他子树中与之颜色相同的点构成一条合法路径。

情况 $2$ 可以通过记录一个颜色数组，将访问过的到 $x$ 的路径上没有与之颜色相同的点的节点的颜色进行统计。两棵不同的子树之间的路径匹配时，直接使用颜色数组即可。

接下来就是点分治模板了。每次取重心，统计贡献，删除，递归点分治子树。即可统计所有合法路径。

实现得比较差，常数较大。

#include <bits/stdc++.h>
using namespace std;
struct edge
{
	long long v,nxt;
}e[600000];
long long t,n,a[300000],s[300000],h[300000],del[300000],he=0,cnt=0,ans=1e10,tol=0;
long long pre[300000],x[300000],sum[300000],g[300000];
void init()
{
	for(int i=1;i<=n;i++)del[i]=0,h[i]=0;
	tol=0,cnt=0;
}

void add_edge(long long u,long long v)
{
	e[++cnt].nxt=h[u];
	e[cnt].v=v;
	h[u]=cnt;
}
 
long long dfs1(long long now,long long fa,long long cnt)
{
    long long maxn=0;
    s[now]=1;
    if(del[now])return 0;
    for(long long i=h[now];i;i=e[i].nxt)
        if(e[i].v!=fa)
            {
            long long z=dfs1(e[i].v,now,cnt);
            s[now]+=z,maxn=max(maxn,z);
            }
    if(max(maxn,cnt-s[now])<ans)ans=min(ans,max(maxn,cnt-s[now])),he=now;
    return s[now];
}

void dfs2(long long now,long long fa)
{
	if(del[now])return;
	s[now]=1,pre[now]=x[a[now]],x[a[now]]=now;
	if(pre[now]==0)
	   {
	   	if(a[now]==a[he])tol++;
	   	else 
	   	   {
	   	   tol+=sum[a[now]];
		   g[a[now]]++;
		   }
	   }
	for(long long i=h[now];i;i=e[i].nxt)
	    if(e[i].v!=fa)
	       {
		   dfs2(e[i].v,now);
		   if(del[e[i].v]==0)s[now]+=s[e[i].v];
	       }
	x[a[now]]=pre[now];
}

void update(long long now,long long fa)
{
	if(del[now])return;
	sum[a[now]]+=g[a[now]],g[a[now]]=0;
	for(long long i=h[now];i;i=e[i].nxt)
	    if(e[i].v!=fa)update(e[i].v,now);
}

void clear(long long now,long long fa)
{
	if(del[now])return;
	pre[now]=0,x[a[now]]=0,sum[a[now]]=0;
	for(long long i=h[now];i;i=e[i].nxt)
	    if(e[i].v!=fa)clear(e[i].v,now);
}
 
void dfz(long long now,long long siz)
{
	if(del[now])return;
	ans=1e10;
	dfs1(now,0,siz);
	for(long long i=h[he];i;i=e[i].nxt)
	    {
	    dfs2(e[i].v,he);
	    update(e[i].v,he);
	    }
	for(long long i=h[he];i;i=e[i].nxt)clear(e[i].v,he);
	del[he]=1;
	for(long long i=h[he];i;i=e[i].nxt)dfz(e[i].v,s[e[i].v]);    
}
 
int main()
{
	scanf("%lld",&t);
	while(t--)
		{
	    scanf("%lld",&n);
	    init();
	    for(long long i=1;i<=n;i++)scanf("%lld",&a[i]);
	    for(long long i=1;i<=n-1;i++)
	        {
	        	long long u=0,v=0;
	        	scanf("%lld%lld",&u,&v);
	        	add_edge(u,v),add_edge(v,u);
			}
		dfz(1,n);
	    printf("%lld\n",tol);
	    }
	return 0;
}