hdu 4670 树的分治-求点对的个数

/*
树的分治
因为树的点权值可达到10^15，注意手动扩栈，还有int64
题意:给你一棵树，给你一些素数，给你每个点一个权值且每个权值均可由这些素数组成。现在定义任意任意两点的价值为他们路径上的权值相乘。求这样的点对的权值为立方数的个数
解：
如果直接求得话会超int64，不可行
由立方数的性质可得，一个数可有素数组成，对于这些素数可以分解为这些素数相乘的形式如，24=(2^3)*(3^1)；如果是立方数的话那么他的各进制对3取余都为0.股24可写成01这种三进制形式
对于这些权值的乘法可有三进制想加可得。
接下来就是树的分治了
当然这里可以先求出一条子树上的各个点的权值乘积，然后和根节点和其他字树比较看是否可以互补那么就找到一对
可用map容器实现。因为他重点是比较到根节点和其他子树是否可以互补，进而递归下去，求出每个子树的这样的点对。
*/
#pragma comment(linker, "/STACK:1024000000,1024000000")
#include<stdio.h>
#include<string.h>
#include<algorithm>
#include<iostream>
#include<map>
#define ll  __int64
#define N 51000
#define inf 0x3fffffff
using namespace std;
struct node {
 ll u,v,next;
}bian[N*4];
ll yong,head[N];
ll num[N],nn,all[N][51],ff[N][51],minn,m,vis[N],fp[N],ma;
ll prime[51],len;
map<ll,ll>q;
void init() {
 yong=0;
 memset(head,-1,sizeof(head));
}
void addedge(ll u,ll v) {
bian[yong].u=u;
bian[yong].v=v;
bian[yong].next=head[u];
head[u]=yong++;
}
void dfs1(ll u,ll fa) {//找到每次根子树的的节点数目
    ll i;
    nn++;
    for(i=head[u];i!=-1;i=bian[i].next) {
        ll v=bian[i].v;
        if(v!=fa&&!vis[v])
            dfs1(v,u);
    }
    return ;
}
ll Max(ll v,ll vv) {
return v>vv?v:vv;
}
void dfs2(ll u,ll fa) {//找重心
 num[u]=1;
 ll i,tit=0;
 for(i=head[u];i!=-1;i=bian[i].next) {
    ll v=bian[i].v;
    if(v!=fa&&!vis[v]) {
        dfs2(v,u);
        num[u]+=num[v];
        tit=Max(tit,num[v]);
    }
 }
 tit=Max(tit,nn-num[u]);
 if(tit<minn) {
    minn=tit;
    ma=u;
 }
 return ;
}
void dfs4(ll u,ll fa) {
 ll i;
 //prllf("dfs4=%d",u);
 if(fa!=-1) {
        ll  e=fp[fa];
    for(i=1;i<=m;i++)
        all[len][i]=(all[e][i]+ff[u][i])%3;
 }
 else {
    for(i=1;i<=m;i++)
        all[len][i]=ff[u][i];
 }
 fp[u]=len++;
 for(i=head[u];i!=-1;i=bian[i].next) {
    ll v=bian[i].v;
    if(v!=fa&&!vis[v])
        dfs4(v,u);
 }
 return ;
 }
ll dfs3(ll u) {//求以当前子树为根的符合条件的对数
   ll i,ans=0,j,k;
   ll s1,s2;
   s1=0;
   q.clear();
   for(i=1;i<=m;i++)
    s1=s1*3+ff[u][i];
  if(s1==0)ans++;
   //prllf("ma=%d,s1=%d\n",ma,s1);
    q[s1]=1;
    for(i=head[u];i!=-1;i=bian[i].next) {
        ll v=bian[i].v;
        if(vis[v])continue;
        len=0;
        dfs4(v,-1);
     //  prllf("len=%d\n",len);
        for(j=0;j<len;j++) {//进行互补
            s1=0;
            for(k=1;k<=m;k++)
                s1=s1*3+(3-all[j][k])%3;
        //  prllf("zzs1=%d\n",s1);
                ans+=q[s1];
        }
        for(j=0;j<len;j++) {
            s2=0;
            for(k=1;k<=m;k++)
                s2=s2*3+(all[j][k]+ff[u][k])%3;//到根结点的权值，为下次互补做准备，同时避免重复
            q[s2]++;
        }
      //  prllf("aad%d\n",ans);
    }
    return ans;
}
ll dfs(ll u) {
  minn=inf;
  nn=0;
  dfs1(u,-1);
  dfs2(u,-1);//求重心
  vis[ma]=1;
  ll ans=dfs3(ma);
 // prllf("ma=%d\n",ma);
  //prllf("ans=%d\n",ans);
  ll i;
  for(i=head[ma];i!=-1;i=bian[i].next) {
    ll v=bian[i].v;
    if(!vis[v])
        ans+=dfs(v);
  }
  return ans;
}
int main() {
      ll n,i,j,k,kk;
      while(scanf("%I64d",&n)!=EOF) {
       init();
       scanf("%I64d",&m);
       for(i=1;i<=m;i++)
        scanf("%I64d",&prime[i]);
       for(i=1;i<=n;i++) {
        scanf("%I64d",&k);
          for(j=1;j<=m;j++) {
            kk=0;
            while(k%prime[j]==0) {
                k/=prime[j];
                kk++;
                kk%=3;
            }
            ff[i][j]=kk;
          }
       }
       for(i=1;i<n;i++) {
        scanf("%I64d%I64d",&j,&k);
        addedge(j,k);
        addedge(k,j);
       }
       memset(vis,0,sizeof(vis));
       printf("%I64d\n",dfs(1));
      }
return 0;}