埃森哲杯第十六届上海大学程序设计联赛春季赛暨上海高校金马五校赛 B-合约数(dfs序)

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

题目描述

在埃森哲，员工培训是最看重的内容，最近一年，我们投入了 9.41 亿美元用于员工培训和职业发展。截至 2018 财年末，我们会在全球范围内设立 100 所互联课堂，将互动科技与创新内容有机结合起来。按岗培训，按需定制，随时随地，本土化，区域化，虚拟化的培训会让你快速取得成长。小埃希望能通过培训学习更多ACM 相关的知识，他在培训中碰到了这样一个问题，

给定一棵n个节点的树，并且根节点的编号为p，第i个节点有属性值val_i, 定义F(i): 在以i为根的子树中，属性值是val_i的合约数的节点个数。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个整数val₁, val₂,…, val_n,其中1<=val_i<=10000，1<=i<=n.

输出描述:

对于每组数据，输出一行，包含1个整数， 表示对1000000007取模后的结果

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const ll mod=1000000007;
vector<int>p[10004];
bool is[10004],vis[20004];
int tol,n,rt,k[20004],cnt[20004],val[20004],head[20004];
struct node
{
    int to,next;
}rode[40004];
void add(int a,int b)
{
    rode[tol].to=b;
    rode[tol].next=head[a];
    head[a]=tol++;
}
void init()
{
    memset(is,1,sizeof(is));
    for(int i=2;i<=10000;i++)
    {
        if(is[i])
        {
            for(int j=i+i;j<=10000;j+=i)
                is[j]=0;
        }
    }
    for(int i=4;i<=10000;i++)
    {
        if(!is[i])
        {
            for(int j=i;j<=10000;j+=i)p[j].push_back(i);//p[j]中保存j的合约数
        }
    }
}
void dfs(int v,int pre)
{
    vis[v]=1;
    for(int i=0;i<p[val[v]].size();i++)
    {
        int x=p[val[v]][i];
        k[v]-=cnt[x];//不在v的子树中的先减去//cnt[x]表示当前每个val出现的次数
    }
    cnt[val[v]]++;
    for(int i=head[v];i!=-1;i=rode[i].next)
    {
        node e=rode[i];
        if(e.to!=pre&&!vis[e.to])dfs(e.to,v);
    }
    for(int i=0;i<p[val[v]].size();i++)
    {
        int x=p[val[v]][i];
        k[v]+=cnt[x];//前面减这里加，抵消掉正好是v的子树中每个x出现的次数
    }
}
int main()
{
    init();
    int T;scanf("%d",&T);
    while(T--)
    {
        memset(head,-1,sizeof(head));
        memset(vis,0,sizeof(vis));tol=1;
        memset(k,0,sizeof(k));
        memset(cnt,0,sizeof(cnt));
        scanf("%d%d",&n,&rt);
        for(int i=1;i<n;i++)
        {
            int x,y;scanf("%d%d",&x,&y);
            add(x,y);add(y,x);
        }
        for(int i=1;i<=n;i++)scanf("%d",&val[i]);
        dfs(rt,0);//根节点开始遍历
        ll ans=0;
        for(ll i=1;i<=n;i++)
            ans=(ans+i*k[i])%mod;
        printf("%lld\n",ans);
    }
    return 0;
}