HDU 4812 D Tree [树分治]

Description

There is a skyscraping tree standing on the playground of Nanjing University of Science and Technology. On each branch of the tree is an integer (The tree can be treated as a connected graph with N vertices, while each branch can be treated as a vertex). Today the students under the tree are considering a problem: Can we find such a chain on the tree so that the multiplication of all integers on the chain (mod 10  ⁶ + 3) equals to K? 
Can you help them in solving this problem? 

题意：给出一棵树，点有权，问是否存在路径使其路径上的权值和MOD 1e6+3为K，若有输出字典序最小的点对<u，v>，否则输出No soluction

解法：很经典的树分治，分治后，计算经过当前分治重心的路径数，计算出每个节点到重心的权值和%K的值（不包括重心），记cot[X]为%K=X的出现次数，minidx[X]为%K=X中所有点的最小下标（保证最小字典序），然后就可以根据其中一个点的值，寻找在其他儿子子树中是否有需要的值即可，然后与当前答案比较更新即可。为了保证是简单路径，在遍历的时候，保证其所属的重心儿子编号U>V即可（详见代码）。

（此题取模数比较小，可以直接开100W的数组来存cot和minidx数组，如果取模数很大，则需要hash，因为N范围通常不会很大，hash复杂度几乎为O（1））

代码：O（NlogN）

#include<stdio.h>
#include<string.h>
#include<algorithm>
#include<math.h>
#include<iostream>
#include<stdlib.h>
#include<set>
#include<map>
#include<queue>
#include<vector>
#include<bitset>
#pragma comment(linker, "/STACK:1024000000,1024000000")
template <class T>
bool scanff(T &ret){ //Faster Input
    char c; int sgn; T bit=0.1;
    if(c=getchar(),c==EOF) return 0;
    while(c!='-'&&c!='.'&&(c<'0'||c>'9')) c=getchar();
    sgn=(c=='-')?-1:1;
    ret=(c=='-')?0:(c-'0');
    while(c=getchar(),c>='0'&&c<='9') ret=ret*10+(c-'0');
    if(c==' '||c=='\n'){ ret*=sgn; return 1; }
    while(c=getchar(),c>='0'&&c<='9') ret+=(c-'0')*bit,bit/=10;
    ret*=sgn;
    return 1;
}
#define inf 1073741823
#define llinf 4611686018427387903LL
#define PI acos(-1.0)
#define lth (th<<1)
#define rth (th<<1|1)
#define rep(i,a,b) for(int i=int(a);i<=int(b);i++)
#define drep(i,a,b) for(int i=int(a);i>=int(b);i--)
#define gson(i,root) for(int i=ptx[root];~i;i=ed[i].next)
#define tdata int testnum;scanff(testnum);for(int cas=1;cas<=testnum;cas++)
#define mem(x,val) memset(x,val,sizeof(x))
#define mkp(a,b) make_pair(a,b)
#define findx(x) lower_bound(b+1,b+1+bn,x)-b
#define pb(x) push_back(x)
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;

#define NN 100100
int n,k,maxv[NN],maxval,center,ptx[NN],lnum,a[NN],sz[NN];
pair<int,int> val[NN],ans;
int vn,d[NN],dn,cot[NN*10],minidx[NN*10];
int mod=1e6+3,inv[NN*10];
bool vis[NN];

struct edge{
    int v,next;
    edge(){}
    edge(int v,int next){
        this->v=v;
        this->next=next;
    }
}ed[NN*2];

void addline(int x,int y){
    ed[lnum]=edge(y,ptx[x]);
    ptx[x]=lnum++;
}

int getsize(int x,int fa){
    sz[x]=1;
    gson(i,x){
        int y=ed[i].v;
        if(vis[y]||y==fa)continue;
        sz[x]+=getsize(y,x);
    }
    return sz[x];
}

void getcenter(int r,int x,int fa){
    maxv[x]=sz[r]-sz[x];
    gson(i,x){
        int y=ed[i].v;
        if(vis[y]||y==fa)continue;
        getcenter(r,y,x);
        maxv[x]=max(maxv[x],sz[y]);
    }
    if(maxv[x]<maxval)maxval=maxv[x],center=x;
}

void getval(int x,int fa,int v){
    val[++vn]=mkp(v,x);
    d[++dn]=v;
    gson(i,x){
        int y=ed[i].v;
        if(vis[y]||y==fa)continue;
        getval(y,x,(ll(v)*a[y])%mod);
    }
}

void calc(int x){
    dn=0;
    minidx[1]=x;
    gson(i,x){
        vn=0;
        int y=ed[i].v;
        if(vis[y])continue;
        getval(y,x,a[y]);
        rep(i,1,vn){
            int t=((ll)k*inv[((ll)val[i].first*a[x])%mod])%mod;
            if(cot[t]){
                int u=minidx[t];
                int v=val[i].second;
                ans=min(ans,min(mkp(u,v),mkp(v,u)));
            }
        }
        rep(i,1,vn){
            int t=val[i].first;
            if(++cot[t]==1)minidx[t]=val[i].second;
            else minidx[t]=min(minidx[t],val[i].second);
        }
    }
    rep(i,1,dn)cot[d[i]]--;
}

void solve(int x){
    maxval=inf;
    getsize(x,0);
    getcenter(x,x,0);
    x=center;
    vis[x]=1;
    calc(x);
    gson(i,x){
        int y=ed[i].v;
        if(vis[y])continue;
        solve(y);
    }
}

void init(int n){
    lnum=0;cot[1]=1;
    ans=mkp(inf,inf);
    rep(i,1,n)ptx[i]=-1,vis[i]=0;
}

int main(){
    inv[1]=1;
    rep(i,2,mod-1)inv[i]=(mod-((ll)(mod/i)*inv[mod%i])%mod)%mod;
    while(scanf("%d%d",&n,&k)!=EOF){
        init(n);
        rep(i,1,n)scanff(a[i]),a[i]%=mod;
        rep(i,1,n-1){
            int x,y;
            scanff(x);scanff(y);
            addline(x,y);addline(y,x);
        }
        solve(1);
        if(ans!=mkp(inf,inf))printf("%d %d\n",ans.first,ans.second);
        else printf("No solution\n");
    }
}