第十六届上海大学程序设计联赛春季赛暨上海高校金马五校赛 B 合约数（启发式合并）

链接：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取模后的结果.

解题思路：深搜，枚举每一颗子树，对每一个子树用map维护val的个数，再判断合约数有多少个，然后计算答案即可。深搜时，用启发式合并（小的合并到大的）来维护map，对于所有合约数，先预处理出来。

#include<iostream>
#include<vector>
#include<string.h>
#include<set>
#include<unordered_map>
using namespace std;
typedef long long int ll;
const int MOD=1e9+7;
struct edge{
    int u;
    int v;
    int next;
}e[40005];
int head[20005],edge_num;
void insert_edge(int u,int v){
    e[edge_num].u=u;
    e[edge_num].v=v;
    e[edge_num].next=head[u];
    head[u]=edge_num++;
}
 
int val[20005];
 
ll ans=0;
vector <int> fac[10005];
 
bool is_prime[20005];
void preprocess(int n)
{
    for (int i=2;i<=n;i++)
        is_prime[i]=true;
    for (int i=2;i<=n;i++)
        if (is_prime[i])
            for (int j=i+i;j<=n;j+=i)
                is_prime[j]=false;
    for (int i=2;i<=n;i++)
        if (!is_prime[i])
            for (int j=i;j<=n;j+=i)
                fac[j].push_back(i);
}
 
unordered_map<int,int> mp[20005];
int mn[20005];
 
int merge(int u,int v){
    if(mp[u].size()<mp[v].size()){
        for(auto &j:mp[u]){
                mp[v][j.first]+=j.second;
        }
        return v;
    }
    else{
        for(auto &j:mp[v]){
                mp[u][j.first]+=j.second;
        }
        return u;
    }
}
 
void dfs(int u,int pre){
    for(int i=head[u];i!=-1;i=e[i].next){
        int v=e[i].v;
        if(v!=pre){
            dfs(v,u);
            mn[u]=merge(mn[u],mn[v]);
        }
    }
 
    mp[mn[u]][val[u]]++;
    ll num=0;
    for(auto &i:fac[val[u]])
        num+=mp[mn[u]][i];
    ans=(ans+num*u%MOD)%MOD;
}
 
 
int main(){
    preprocess(10000);
    int T;
    int u,v;
    int N,P;
    scanf("%d",&T);
    while(T--){
        for(int i=0;i<=20002;i++){
            mn[i]=i;
            mp[i].clear();
        }
        ans=0;
        memset(head,-1,sizeof(head));
        edge_num=0;
        scanf("%d%d",&N,&P);
        for(int i=0;i<N-1;i++){
            scanf("%d%d",&u,&v);
            insert_edge(u,v);
            insert_edge(v,u);
        }
        for(int i=1;i<=N;i++)
            scanf("%d",&val[i]);
 
        dfs(P,-1);
        printf("%lld\n",ans);
 
    }
 
    return 0;
}