hdu4670 树分治

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <vector>
#include <string.h>
#include <map>
using namespace std;
const int maxn=50005;
typedef long long LL;
int H[maxn],nx[maxn*2],to[maxn*2],numofE;
LL prime[40],V[maxn],ans;
LL bas[40];
int N,K;
LL fenjie(LL a)
{
    LL an=0;
    for(int i=0; i<K; i++)
    {
         int d=0;
         while((a%prime[i])==0){
            d++;a/=prime[i];
         }
         an+=bas[i]*(d%3);
    }
    return an;
}
void add(int u, int v)
{
    numofE++;
    to[numofE]=v;
    nx[numofE]=H[u];
    H[u]=numofE;
}
int Q[maxn],fa[maxn],subnum[maxn];
LL P[maxn];
bool center[maxn];
int searchroot(int cur)
{
    int rear=0;
    fa[cur]=-1,Q[rear++]=cur;
    for(int i=0; i<rear; i++)
    {
        int x=Q[i];
        for(int j = H[x]; j; j=nx[j])
        {
            int tto = to[j];
            if( tto == fa[x] || center[tto] )continue;
            Q[rear++]=tto; fa[tto]=x;
        }
    }
    int MIN=maxn*2,root=cur;
    for(int i=rear-1; i>=0; i--)
    {
        int x=Q[i];
        subnum[x]=1;
        int MA=0;
        for(int j=H[x]; j; j=nx[j])
        {
             int tto=to[j];
             if(tto == fa[x] || center[tto] )continue;
             MA=max(MA,subnum[tto]);
             subnum[x]+=subnum[tto];
        }
        MA=max(MA,rear-subnum[x]);
        if(MIN>MA){
            MIN=MA; root=x;
        }
    }
    return root;
}
LL requrenum(LL a)
{
    LL an=0;
    for(int i=0; i<K; i++)
    {
        LL d=3-(a%3);
        a/=3;
        if(d>2)d=0;
        an+=d*bas[i];
    }
    return an;
}
LL temp[maxn];
LL jia(LL a, LL b)
{
    LL an=0;
    for(int i=0; i<K; i++)
    {
        int d=(a%3 + b%3)%3;
        a/=3;b/=3;
        an+=bas[i]*d;
    }
    return an;
}
void count_pair(map<LL,LL>&ds,map<LL,LL>tds,int root)
{
   map<LL,LL>::iterator it;
   it=tds.begin();
   while(it!=tds.end()){
     LL re=requrenum(it->first);
     if(ds.count(re)){
        ans+=ds[re]*(it->second);
     }
     ++it;
   }
   it=tds.begin();
   while(it!=tds.end()){
     LL vv=jia(it->first,V[root]);
     if(ds.count(vv)){
        ds[vv]+=it->second;
     }else
        ds[vv]=it->second;
     it++;
   }
}
void updateedg(int cur, map<LL,LL> &ds)
{
    int rear=0;
    fa[cur]=0;
    Q[rear++]=cur;
    P[0]=V[cur];
    if(ds.count(P[0]))ds[P[0]]++;
    else ds[P[0]]=1;
    for(int i=0; i<rear; i++)
    {
        int x=Q[i];
        for(int j=H[x]; j; j=nx[j])
        {
            int tto=to[j];
            if(tto==fa[x]||center[tto])continue;
            fa[tto]=x;
            P[rear]=jia(P[i],V[tto]);
            if(ds.count(P[rear]))ds[ P[rear] ]++;
            else ds[ P[rear] ]=1;
            Q[rear++]=tto;
        }
    }
}
void dfs(int cur)
{
    int root;
    root=searchroot(cur);
    center[root]=true;
    map<LL,LL>ds,tds;
    ds[V[root]]=1;
    for(int i=H[root]; i; i=nx[i])
    {
        int tto=to[i];
        if(center[tto])continue;
        dfs(tto);
        tds.clear();
        updateedg(tto,tds);
        count_pair(ds,tds,root);
    }
    center[root]=false;

    return ;
}
int main()
{
    bas[0]=1;
    for(int i=1;i<=32; i++)
        bas[i]=bas[i-1]*3;
    while(scanf("%d",&N)==1)
        {
               scanf("%d",&K);
                numofE=ans=0;
               for(int i=0; i<K; i++)scanf("%I64d",&prime[i]);
               for(int i=1; i<=N; i++){
                    H[i]=0;
                    scanf("%I64d",&V[i]);
                    V[i]=fenjie(V[i]);
                    if(V[i]==0)ans++;
               }
               for(int i=1; i<N; i++)
               {
                   int a,b;
                   scanf("%d%d",&a,&b);
                   add(a,b);
                   add(b,a);
               }
               dfs(1);
               printf("%I64d\n",ans);
        }
    return 0;
}