牛客-合约数（树上启发式合并）

题目描述：点击进入

题意

给定一棵 n 个节点的树，并且根节点的编号为 p ，第 i 个节点有属性值 val（ i ）, 定义 F ( i ) : 在以 i 为根的子树中，属性值是 val ( i ) 的合约数的节点个数。
y 是 x 的合约数是指 y 是合数且 y 是 x 的约数。

思路

只对子树的查询且没有修改，这不就是树上启发式合并嘛。。。
要求节点 i 的属性值（ val ( i ) ） 在以 i 为根节点的子树中合约数的个数
我们可以提前预处理出所有数的合约数（ val 最大到 1e4 ）
统计的时候可以直接枚举每个子树根节点属性值的合约数，
将他们在子树中出现的次数相加就是根节点在子树中合约数的个数

代码

#include<iostream>
#include<string>
#include<map>
#include<set>
#include<queue>
#include<cstdio>
#include<vector>
#include<cstring>
#include<algorithm>
#include<iomanip>
#include<cmath>
#include<fstream>
#define X first
#define Y second
#define INF 0x3f3f3f3f
#define pii pair<int, int>
using namespace std;
typedef long long ll;
typedef unsigned long long llu; 
const int maxn=2e4+10;
const int mod=1000000007;
int t,n,x,y,p,val[maxn],son[maxn],cnt[maxn],fson,head[maxn],tot,sum[maxn];
ll ans[maxn];
bool vis[maxn];
vector<int>v[maxn];
struct node
{
	int to;
	int next;
}edge[maxn<<1];
void noprime()//求每个数的合约数
{
	vis[1]=0;
	for(int i=2;i<=10000;i++)//线性筛求合数
	{
		if(!vis[i])
		{
			for(int j=i+i;j<=10000;j+=i)
			vis[j]=1;
		}
	}
	for(int i=1;i<=10000;i++)//求合约数
	{
		for(int j=1;j<=sqrt(i);j++)
		{
			if(i%j==0)
			{
				if(vis[j]) 
				    v[i].push_back(j);
				if(vis[i/j]&&j*j!=i) 
				    v[i].push_back(i/j);
			}
		} 
	}
}
void init()
{
	memset(head,-1,sizeof(head));
	memset(ans,0,sizeof(ans));
	memset(cnt,0,sizeof(cnt));
	memset(son,0,sizeof(son));
	memset(vis,0,sizeof(vis));
	memset(sum,0,sizeof(sum));
	tot=0;
}
void add(int u,int v)
{
	edge[tot].to=v;
	edge[tot].next=head[u];
	head[u]=tot++;
}
void dfs(int pos,int fa)
{
	sum[pos]=1;
	for(int i=head[pos];i!=-1;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v==fa) continue;
		dfs(v,pos);
		sum[pos]+=sum[v];
		if(sum[son[pos]]<sum[v]) son[pos]=v;
	}
}
void change(int pos,int fa,int d)
{
	cnt[val[pos]]+=d;
	for(int i=head[pos];i!=-1;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v==fa||v==fson) continue;
		change(v,pos,d);
	}
}
void dsu(int pos,int fa,int keep)
{
	for(int i=head[pos];i!=-1;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v==fa||v==son[pos]) continue;
		dsu(v,pos,0);
	}
	if(son[pos])
	{
		dsu(son[pos],pos,1);
		fson=son[pos];
	}
	change(pos,fa,1);
	fson=0;
	for(int i=0;i<v[val[pos]].size();i++)
	ans[pos]+=cnt[v[val[pos]][i]];
	//枚举这个节点属性值的所有合约数，将他们在子树中出现的次数相加
	if(!keep) change(pos,fa,-1);
}
int main( )
{
	ios::sync_with_stdio(false);
	noprime();
	cin>>t;
	while(t--)
	{
		init();
		cin>>n>>p;
		for(int i=1;i<n;i++)
		{
			cin>>x>>y;
			add(x,y);
	    	add(y,x);
		}
		for(int i=1;i<=n;i++) cin>>val[i];
		dfs(p,0);
		dsu(p,0,1);
		ll tmp=0;
		for(int i=1;i<=n;i++)
			tmp=(tmp+i*ans[i])%mod;
		cout<<tmp<<endl;
	}
    return 0;
}