合约数

题目：
链接：https://www.nowcoder.com/acm/contest/91/B
来源：牛客网

给定一棵n个节点的树，并且根节点的编号为p，第i个节点有属性值vali, 定义F(i): 在以i为根的子树中，属性值是vali的合约数的节点个数。y 是 x 的合约数是指 y 是合数且 y 是 x 的约数。小埃想知道对1000000007取模后的结果.
输入描述:
输入测试组数T，每组数据，输入n+1行整数，第一行为n和p，1<=n<=20000, 1<=p<=n, 接下来n-1行，每行两个整数u和v,表示u和v之间有一条边。第n+1行输入n个整数val1, val2,…, valn,其中1<=vali<=10000，1<=i<=n.
输出描述:
对于每组数据，输出一行，包含1个整数， 表示对1000000007取模后的结果
示例1
输入

2
5 4
5 3
2 5
4 2
1 3
10 4 3 10 5
3 3
1 3
2 1
1 10 1
输出

11
2
备注:
n>=10000的有20组测试数据

合数：除1与本身外，还存在其他的除数
dfs到该节点时还未遍历其子节点再次回溯到该节点时已经遍历过其子节点，两者相减即为子节点中合约数的个数

代码：

#include<bits/stdc++.h> 
using namespace std;
typedef long long ll;
const int N=20005;

vector<int> g[N];
int val[N];
bool isprime[N];
vector<short> f[N];
int cnt[N],F[N];
ll res;
const ll MOD =1E9+7;

void init()
{
    for(int i=2;i<N;i++) isprime[i]=1;
    for(int i=2;i<N;i++)//找出合数
    {
        if(isprime[i])
        {
            for(int j=i+i;j<N;j+=i)
            {
                isprime[j]=0;
             } 
        }
    }
    for(int i=4;i<=10000;i++)//找出每个数的合约数并用f[]进行存储
    {
        if(!isprime[i])
        {
            for(int j=i;j<=10000;j+=i)
            {
                f[j].push_back(i);

            }
        }
    }
    //for(int i=0;i<f[20].size();i++) printf("%d ",f[20][i]);
}

void dfs(int u,int par)//节点和父节点
{
     for(int i=0;i<f[val[u]].size();i++)
     {
        int num=f[val[u]][i];
        F[u]-=cnt[num];//先减去之前出现过的
     }
     cnt[val[u]]++;//cnt统计出现的次数
     for(int i=0;i<g[u].size();i++)
     {
        int v=g[u][i];
        if(v==par) continue;
        dfs(v,u);
     }
     for(int i=0;i<f[val[u]].size();i++)
     {
        int num=f[val[u]][i];
        F[u]+=cnt[num];
     }
     res+=F[u]*1LL*u;
     res%=MOD;

}

int main()
{
    int T;
    init();
    scanf("%d",&T);
    while(T--)
    {
        memset(cnt,0,sizeof(cnt));
        res=0;
        int n,p;
        scanf("%d %d",&n,&p);
        for(int i=1;i<=n;i++) F[i]=0,g[i].clear();
        for(int i=1;i<n;i++)
        {
            int u,v;
            scanf("%d %d",&u,&v);
            g[u].push_back(v);
            g[v].push_back(u);

        }
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&val[i]);
        }
        dfs(p,-1);
        printf("%lld\n",res);
    }
    return 0;
}