HDU4812 D Tree（树的点分治）

  1 #include<cstdio>
  2 #include<cstring>
  3 #include<algorithm>
  4 using namespace std;
  5 #define INF (1<<30)
  6 #define MAXN 111111
  7 struct Edge{
  8     int v,next;
  9 }edge[MAXN<<1];
 10 int NE,head[MAXN];
 11 void addEdge(int u,int v){
 12     edge[NE].v=v; edge[NE].next=head[u]; head[u]=NE++;
 13 }
 14 bool vis[MAXN];
 15 int mini,cen,size[MAXN];
 16 void getSize(int u,int fa){
 17     size[u]=1;
 18     for(int i=head[u]; i!=-1; i=edge[i].next){
 19         int v=edge[i].v;
 20         if(v==fa || vis[v]) continue;
 21         getSize(v,u);
 22         size[u]+=size[v];
 23     }
 24 }
 25 void getCen(int u,int fa,int &tot){
 26     int res=tot-size[u];
 27     for(int i=head[u]; i!=-1; i=edge[i].next){
 28         int v=edge[i].v;
 29         if(v==fa || vis[v]) continue;
 30         getCen(v,u,tot);
 31         res=max(res,size[v]);
 32     }
 33     if(res<mini) mini=res,cen=u;
 34 }
 35 int getCen(int u){
 36     getSize(u,u);
 37     mini=INF;
 38     getCen(u,u,size[u]);
 39     return cen;
 40 }
 41 long long ine(long long a){
 42     long long res=1,n=1000001;
 43     while(n){
 44         if(n&1) res*=a,res%=1000003;
 45         a*=a; a%=1000003;
 46         n>>=1;
 47     }
 48     return res;
 49 }
 50 int n,k,val[MAXN];
 51 int ansx,ansy;
 52 int record[1000003],tn,tmpx[MAXN],tmpy[MAXN],all[MAXN],an;
 53 void dfs(int u,int fa,long long dist,int &top){
 54     int v=record[ine(dist)*k%1000003*top%1000003];
 55     if(v){
 56         if(u<v){
 57             if(u<ansx) ansx=u,ansy=v;
 58             else if(u==ansx && v<ansy) ansy=u,ansy=v;
 59         }else{
 60             if(v<ansx) ansx=v,ansy=u;
 61             else if(v==ansx && u<ansy) ansy=v,ansy=u;
 62         }
 63     }
 64     tmpx[tn]=u; tmpy[tn]=dist; ++tn;
 65     all[an++]=dist;
 66     for(int i=head[u]; i!=-1; i=edge[i].next){
 67         int v=edge[i].v;
 68         if(v==fa || vis[v]) continue;
 69         dfs(v,u,dist*val[v]%1000003,top);
 70     }
 71 }
 72 void conquer(int u){
 73     an=0;
 74     all[an++]=val[u];
 75     record[val[u]]=u;
 76     for(int i=head[u]; i!=-1; i=edge[i].next){
 77         int v=edge[i].v;
 78         if(vis[v]) continue;
 79         tn=0;
 80         dfs(v,v,(long long)val[u]*val[v]%1000003,val[u]);
 81         for(int j=0; j<tn; ++j){
 82             if(record[tmpy[j]]==0 || record[tmpy[j]]>tmpx[j]) record[tmpy[j]]=tmpx[j];
 83         }
 84     }
 85     for(int i=0; i<an; ++i) record[all[i]]=0;
 86 }
 87 void divide(int u){
 88     u=getCen(u);
 89     vis[u]=1;
 90     conquer(u);
 91     for(int i=head[u]; i!=-1; i=edge[i].next){
 92         int v=edge[i].v;
 93         if(vis[v]) continue;
 94         divide(v);
 95     }
 96 }
 97 int main(){
 98     int a,b;
 99     while(~scanf("%d%d",&n,&k)){
100         for(int i=1; i<=n; ++i){
101             scanf("%d",val+i);
102         }
103         NE=0;
104         memset(head,-1,sizeof(head));
105         for(int i=1; i<n; ++i){
106             scanf("%d%d",&a,&b);
107             addEdge(a,b);
108             addEdge(b,a);
109         }
110         memset(vis,0,sizeof(vis));
111         ansx=ansy=INF;
112         divide(1);
113         if(ansx==INF) puts("No solution");
114         else printf("%d %d\n",ansx,ansy);
115     }
116     return 0;
117 }