【金马五校赛】B 合约数（预处理+dfs序）

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

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

给定一棵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组测试数据

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

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

给定一棵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组测试数据

题解：

代码：

#include <bits/stdc++.h>
using namespace std;

#define rep(i,a,n) for(int i=a;i<n;i++)
#define mem(a,n) memset(a,n,sizeof(a))
#define pb push_back
#define IO ios::sync_with_stdio(false);
typedef long long ll;
const int N=2e4+5;
const double eps=1e-4;
const int MOD=1e9+7;
const int dir[4][2]= {-1,0,1,0,0,-1,0,1}; ///上下左右
const int INF=0x3f3f3f3f;

vector<int>fac[N];
vector<int>g[N];
bool isprime[N];
int val[N],cnt[N];
int F[N];
///预处理 [1,N]中每个数的合约数个数  O(N*sqrt(N))
void init() 
{
    rep(i,2,N) isprime[i]=1;
    for(int i=2; i*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++)
    {
        if(!isprime[i])
        {
            for(int j=i; j<=N; j+=i)
                fac[j].push_back(i);
        }
    }
}
int ans;
void dfs(int u,int fa)
{
    for(auto num:fac[val[u]])
    {
        F[u]-=cnt[num];///先减去之前出现过的合约数的个数
    }
    cnt[val[u]]++;
    for(auto v:g[u])
    {
        if(v==fa) continue;
        dfs(v,u);///递归树
    }
    ///回溯
    for(auto num:fac[val[u]])
    {
        F[u]+=cnt[num];///现在出现过的总数 = 遍历完其子树后的总合约数个数-之前出现的合约数个数.  
    }
    ans=(ans+1LL*u*F[u])%MOD;
}
int main()
{
    init();
    int T;
    scanf("%d",&T);
    while(T--)
    {
        int n,p;
        scanf("%d%d",&n,&p);
        mem(cnt,0);
        for(int i=0; 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]);
        }
        ans=0;
        dfs(p,-1);
        printf("%d\n",ans);
    }
    return 0;
}