【树上启发式合并】牛客：合约数（数学+奇怪的记录方式）

【树上启发式合并】牛客：合约数（数学+奇怪的记录方式）

题意：

给定一棵n个节点的树，并且根节点的编号为p，第i个节点有属性值vali, 定义F(i): 在以i为根的子树中，属性值是vali的合约数的节点个数。y 是 x 的合约数是指 y 是合数且 y 是 x 的约数。小埃想知道${\sum }_{i=1}^{n}i\cdot F\left(i\right)$

思路：

因为所有数据范围$\le 2e^4$，可以预处理出每个数的合约数，用$vector$存储，再树上启发式合并统计

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int inf=0x3f3f3f3f;
const int maxn=2e4+5;
const ll mod=1e9+7;
int tot,prim[maxn],book[maxn],siz[maxn],dfn[maxn],tmp[maxn];
int n,rt,t,num,son[maxn],id[maxn],w[maxn],fd[maxn];
ll ans,sum;
vector<int>mp[maxn],p[maxn];
void init(){
    for(int i=2;i<=maxn;i++){
    	if(!book[i])prim[tot++]=i;
    	for(int j=0;j<tot&&prim[j]*i<maxn;j++){
    		book[i*prim[j]]=true;
    		if(i%prim[j]==0)break;
		}
	} 
	for(int i=1;i<=maxn;i++){
		for(int j=1;j*j<=i;j++){
			if(i%j==0){
				if(book[j])p[i].push_back(j);
				if(j!=i/j&&book[i/j])p[i].push_back(i/j);
			}
		}
	}
}
void dfs1(int x,int fa){
	siz[x]=1;
	dfn[x]=++num;
	id[num]=x;
	for(auto v:mp[x]){
		if (v==fa)continue;
		dfs1(v,x);
		siz[x]+=siz[v];
		if (siz[v]>siz[son[x]])son[x]=v;
	}
}
void dfs(int x,int f,int op){
	for(int i = 0; i < mp[x].size(); i++) {
        int y = mp[x][i];
		if(y!=f&&y!=son[x])dfs(y,x,0);
	}
	if(son[x])dfs(son[x],x,1);
	for(int i = 0; i < mp[x].size(); i++) {
        int v = mp[x][i];
		if(v==f||v==son[x])continue;
		for(int j=dfn[v];j<=dfn[v]+siz[v]-1;j++){
			int y=id[j];
			tmp[w[y]]++;
		}
	}
	tmp[w[x]]++;
	for(auto h:p[w[x]])
		fd[x]+=tmp[h];
	if(!op){
		for(int j=dfn[x];j<=dfn[x]+siz[x]-1;j++){
			int y=id[j];
			tmp[w[y]]--;
		}
	}
}
int main(){
    init();
    scanf("%d",&t);
    while(t--){
    	ans=num=0;
    	scanf("%d%d",&n,&rt);
    	memset(fd,0,sizeof(fd));
    	memset(son,0,sizeof(son));
    	for(int i=1;i<=n;i++)mp[i].clear();
        for(int i=1,x,y;i<n;i++){
    	    scanf("%d%d",&x,&y);
    	    mp[x].push_back(y);
    	    mp[y].push_back(x);
	    }
	    for(int i=1;i<=n;i++)
    	    scanf("%d",&w[i]);
        dfs1(rt,0);
        dfs(rt,0,0);
        for(int i=1;i<=n;i++){
    	    ans+=1ll*i*fd[i];
	    }
        printf("%lld\n",ans%mod);
	}
}