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

题目描述

在埃森哲，员工培训是最看重的内容，最近一年，我们投入了 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取模后的结果

用启发式合并，从集合小的向上合并，用map直接查询

#include <iostream>
#include <algorithm>
#include <cstring>
#include <cstdio>
#include <vector>
#include <queue>
#include <stack>
#include <cstdlib>
#include <iomanip>
#include <cmath>
#include <cassert>
#include <ctime>
#include <map>
#include <set>
using namespace std;
#pragma comment(linker, "/stck:1024000000,1024000000")
#define lowbit(x) (x&(-x))
#define max(x,y) (x>=y?x:y)
#define min(x,y) (x<=y?x:y)
#define MAX 100000000000000000
#define MOD 1000000007
#define pi acos(-1.0)
#define ei exp(1)
#define PI 3.1415926535897932384626433832
#define ios() ios::sync_with_stdio(true)
#define INF 0x3f3f3f3f
#define mem(a) (memset(a,0,sizeof(a)))
#define ll long long
int n,val[20006],p,t,x,y,top;
int head[20006],prime[10006];
vector<int>fac[10005];
map<int,int>m[20005];
map<int,int>::iterator it;
int ans[20006];
ll cnt;
struct node
{
    int u,v,next;
}e[40006];
void add(int u,int v)
{
    e[top].u=u;
    e[top].v=v;
    e[top].next=head[u];
    head[u]=top++;
}
void init()
{
    memset(prime,1,sizeof(prime));
    for(int i=2;i<=10000;i++)
    {
        if(!prime[i]) continue;
        for(int j=2;j*i<=10000;j++)
            prime[j*i]=0;
    }
    for(int i=2;i<=10000;i++)
    {
        if(prime[i]) continue;
        for(int j=i;j<=10000;j+=i)
            fac[j].push_back(i);//求j的所有约数
    }
}
int unio(int u,int v)
{
    if(m[u].size()<m[v].size())
    {
        for(it=m[u].begin();it!=m[u].end();++it)
            m[v][it->first]+=it->second;
        return v;
    }
    else
    {
        for(it=m[v].begin();it!=m[v].end();it++)
            m[u][it->first]+=it->second;
        return u;
    }
}
void dfs(int u,int pre)
{
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int now=e[i].v;
        if(now!=pre)
        {
            dfs(now,u);
            ans[u]=unio(ans[u],ans[now]);
        }
    }
    m[ans[u]][val[u]]++;
    ll num=0;
    for(int i=0;i<fac[val[u]].size();i++)
        num+=m[ans[u]][fac[val[u]][i]];//若有value的约数存在则有值，不存在则为0；
    cnt=(cnt+u*num%MOD)%MOD;
}
int main()
{
    init();
    scanf("%d",&t);
    while(t--)
    {
        scanf("%d%d",&n,&p);
        for(int i=1;i<=n;i++)
        {
            ans[i]=i;
            m[i].clear();
        }
        memset(head,-1,sizeof(head));
        top=0;
        cnt=0;
        for(int i=1;i<n;i++)
        {
            scanf("%d%d",&x,&y);
            add(x,y);
            add(y,x);
        }
        for(int i=1;i<=n;i++)
            scanf("%d",&val[i]);
        dfs(p,-1);
        printf("%lld
",cnt);
    }
    return 0;
}