HDU4670 Cube number on a tree 树分治

    人生的第一道树分治,要是早点学我南京赛就不用那么挫了,树分治的思路其实很简单,就是对子树找到一个重心(Centroid),实现重心分解,然后递归的解决分开后的树的子问题,关键是合并,当要合并跨过重心的两棵子树的时候,需要有一个接近O(n)的方法,因为f(n)=kf(n/k)+O(n)解出来才是O(nlogn).在这个题目里其实就是将第一棵子树的集合里的每个元素,判下有没符合条件的,有就加上,然后将子树集合压进大集合,然后继续搞第二棵乃至第n棵.我的过程用了map,合并是nlogn的所以代码速度颇慢,大概6s,题目时限10s,可以改成hash应该会快许多,毕竟用map实在太慢,用vector也可以,具体可以参见挑战程序设计竞赛代码.下面的代码查找重心用了挑战的代码.

#pragma comment(linker, "/STACK:102400000,102400000")

#include<iostream>

#include<cstring>

#include<string>

#include<cstdio>

#include<algorithm>

#include<map>

#include<vector>

#define maxv 50000

#define ll long long

using namespace std;



int n,k;

vector<int> G[maxv+50];

ll val[maxv+50];

ll prime[maxv+50];

ll convert_three(ll v)

{

    ll bas=1;ll res=0;

    for(int i=0;i<k;++i){

        int num=0;

        while(v%prime[i]==0){

            v/=prime[i];

            num++;

        }

        num%=3;res+=num*bas;

        bas*=3;

    }

    return res;

}



ll xor(ll x,ll y)

{

    ll res=0;ll bas=1;

    for(int i=0;i<k;++i){

        res+=((x%3)+(y%3))%3*bas;

        x/=3;y/=3;

        bas*=3;

    }

    return res;

}



ll inv(ll x)

{

    ll res=0;ll bas=1;

    for(int i=0;i<k;++i){

        res+=((3-(x%3))%3)*bas;

        x/=3;

        bas*=3;

    }

    return res;

}



void print(ll x){

    while(x){

        cout<<x%3;

        x/=3;

    }

    cout<<endl;

}



bool centroid[maxv+50];

int ssize[maxv+50];

int ans;



map<ll,int> sta;

map<ll,int>::iterator it;

int compute_ssize(int v,int p)

{

    int c=1;

    for(int i=0;i<G[v].size();++i){

        int w=G[v][i];

        if(w==p||centroid[w]) continue;

        c+=compute_ssize(G[v][i],v);

    }

    ssize[v]=c;

    return c;

}



pair<int,int> search_centroid(int v,int p,int t)

{

    pair<int,int> res=make_pair(INT_MAX,-1);

    int s=1,m=0;

    for(int i=0;i<G[v].size();++i){

        int w=G[v][i];

        if(w==p||centroid[w]) continue;

        res=min(res,search_centroid(w,v,t));

        m=max(m,ssize[w]);

        s+=ssize[w];

    }

    m=max(m,t-s);

    res=min(res,make_pair(m,v));

    return res;

}



void enumerate_mul(int v,int p,ll d,map<ll,int> &ds)

{

    if(ds.count(d)) ds[d]++;

    else ds[d]=1;

    for(int i=0;i<G[v].size();++i){

        int w=G[v][i];

        if(w==p||centroid[w]) continue;

        enumerate_mul(w,v,xor(d,val[w]),ds);

    }

}



void solve(int v)

{

    compute_ssize(v,-1);

    int s=search_centroid(v,-1,ssize[v]).second;

    centroid[s]=true;

    for(int i=0;i<G[s].size();++i){

        if(centroid[G[s][i]]) continue;

        solve(G[s][i]);

    }

    sta.clear();

    sta[val[s]]=1;map<ll,int> tds;

    for(int i=0;i<G[s].size();++i){

        if(centroid[G[s][i]]) continue;

        tds.clear();

        enumerate_mul(G[s][i],s,val[G[s][i]],tds);

        it=tds.begin();

        while(it!=tds.end()){

            ll rev=inv((*it).first);

            if(sta.count(rev)){

                ans+=sta[rev]*(*it).second;

            }

            ++it;

        }

        it=tds.begin();

        while(it!=tds.end()){

            ll  vv=xor((*it).first,val[s]);

            if(sta.count(vv)){

                sta[vv]+=(*it).second;

            }

            else{

                sta[vv]=(*it).second;

            }

            ++it;

        }

    }

    centroid[s]=false;

}



int main()

{

    while(cin>>n>>k){

        ans=0;

        for(int i=0;i<k;++i){

            scanf("%I64d",&prime[i]);

        }

        G[0].clear();

        for(int i=1;i<=n;++i){

            scanf("%I64d",&val[i]);

            val[i]=convert_three(val[i]);

            if(val[i]==0) ans++;

            //print(val[i]);

            G[i].clear();

        }

        int u,v;

        for(int i=0;i<n-1;++i){

            scanf("%d%d",&u,&v);

            G[u].push_back(v);

            G[v].push_back(u);

        }

        memset(centroid,0,sizeof(centroid));

        solve(1);

        printf("%d\n",ans);

    }

    return 0;

}