计蒜客 16初赛 C 树上路径+容斥

题意:n个点的树,每个点的价值为a[i],
n<=1e4,a[i]<=500. 求所有(a[i],a[j])互质的点的距离和


树上路径求和的问题,从边的角度来考虑.计算u-v这条边总共走了多少次.
设dp[v][j] 子树v中,值为j的个数, 可以枚举j = [1:500] 若子树v上方有num个数和j互质,那么贡献 num*dp[v][j]
直接枚举的话 O(n*[ai]^2) TLE...

因为a[i]<=500 那么每个数最多只有4个素因子,考虑gcd(j,x)!=1的个数.

算出子树v上方gcd为p[j]的倍数有多少个. 容斥搞一搞即可.O(n*a[i]*k*2^k) k<=3)

#include <bits/stdc++.h>
using namespace std;
const int N=2e4+5,M=505;
int n,a[N];
vector<int> p[M],e[N];
void init()
{
	for(int i=1;i<=500;i++)
	{
		int x=i;
		for(int j=2;j*j<=i;j++)
		{
			if(x%j==0)
			{
				p[i].push_back(j);
				while(x%j==0)
					x/=j;
			}
		}
		if(x!=1)
			p[i].push_back(x);
	}
}
int dp[N][M],res=0;//dp[u][j] 子树u中 值为j的个数.  
int cnt[M];//cnt[x] 为素因子x倍数的顶点数 
int sum[N][M];//sum[u][j] 子树u中 为j倍数的个数. 
void dfs(int u,int fa)
{
	int x=a[u];
	dp[u][x]++;
	for(int i=0;i<(1<<p[x].size());i++)
	{
		int num=1;
		for(int j=0;j<p[x].size();j++)
		{
			if((i>>j)&1)
				num*=p[x][j];
		}
		sum[u][num]++;
	}
	for(int i=0;i<e[u].size();i++)
	{
		int v=e[u][i];
		if(v==fa)	continue;
		dfs(v,u);
		for(int j=1;j<=500;j++)
		{
			dp[u][j]+=dp[v][j];
			sum[u][j]+=sum[v][j];
		}
		//u-v 这条边贡献: dp[v][j] * num(v上半部分和j互质的个数) 
		for(int j=1;j<=500;j++)
		{
			for(int s=0;s<(1<<p[j].size());s++)
			{
				int ct=0,num=1;
				for(int k=0;k<p[j].size();k++)
				{
					if((s>>k)&1)
						ct++,num*=p[j][k];
				}
				if(ct%2)	res-=dp[v][j]*(cnt[num]-sum[v][num]);
				else		res+=dp[v][j]*(cnt[num]-sum[v][num]);
			}	
		} 
	}
}
int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	init();
	cin>>n;
	for(int i=1;i<=n;i++)
	{
		cin>>a[i];
		int x=a[i];
		for(int s=0;s<(1<<p[x].size());s++)
		{
			int num=1;
			for(int k=0;k<p[x].size();k++)
			{
				if((s>>k)&1)
					num*=p[x][k];
			}
			cnt[num]++; 
		}
	}
	for(int i=1;i<=n-1;i++)
	{
		int u,v;
		cin>>u>>v;
		e[u].push_back(v);
		e[v].push_back(u);
	}
	dfs(1,0);
	cout<<res<<'\n';
	return 0;
}