HDU 4812 D Tree

  1 #include <bits/stdc++.h>
  2 #define A first
  3 #define B second
  4 using namespace std;
  5 using PII = pair<int,int>;
  6 using LL = long long;
  7 const int B = 100000;
  8 const int maxn = 100010;
  9 const int mod = 1e6 + 3;
 10 const PII INF = PII(~0U>>2,~0U>>2);
 11 struct arc {
 12     int to,next;
 13     arc(int x = 0,int y = -1) {
 14         to = x;
 15         next = y;
 16     }
 17 } e[maxn<<1];
 18 bool vis[maxn];
 19 LL inv[1000010],val[maxn],F[1000010],K,P;
 20 PII ans;
 21 int head[maxn],sz[maxn],maxson[maxn],tot,n;
 22 void init() {
 23     inv[1] = 1;
 24     for(int i = 2; i < mod; ++i)
 25         inv[i] = (mod - mod/i)*inv[mod%i]%mod;
 26 }
 27 void add(int u,int v) {
 28     e[tot] = arc(v,head[u]);
 29     head[u] = tot++;
 30 }
 31 void dfs(int u,int fa) {
 32     sz[u] = 1;
 33     maxson[u] = 0;
 34     for(int i = head[u]; ~i; i = e[i].next) {
 35         if(e[i].to == fa || vis[e[i].to]) continue;
 36         dfs(e[i].to,u);
 37         sz[u] += sz[e[i].to];
 38         maxson[u] = max(maxson[u],sz[e[i].to]);
 39     }
 40 }
 41 int FindRoot(int sum,int u,int fa) {
 42     int ret = u;
 43     maxson[u] = max(maxson[u],sum - sz[u]);
 44     for(int i = head[u]; ~i; i = e[i].next) {
 45         if(e[i].to == fa || vis[e[i].to]) continue;
 46         int x = FindRoot(sum,e[i].to,u);
 47         if(maxson[x] < maxson[ret]) ret = x;
 48     }
 49     return ret;
 50 }
 51 void calc(int u,int fa,LL d,int flag){
 52     d = d*val[u]%mod;
 53     if(flag){
 54         if(F[K*inv[d]%mod] > P){
 55             int x = F[K*inv[d]%mod] - P,y = u;
 56             if(x > y) swap(x,y);
 57             ans = min(ans,PII(x,y));
 58         }
 59     }else if(F[d] <= P) F[d] = P + u;
 60     else F[d] = min(F[d],P + u);
 61     for(int i = head[u]; ~i; i = e[i].next){
 62         if(e[i].to == fa || vis[e[i].to]) continue;
 63         calc(e[i].to,u,d,flag);
 64     }
 65 }
 66 void solve(int u) {
 67     dfs(u,0);
 68     int rt = FindRoot(sz[u],u,0);
 69     vis[rt] = true;
 70     P += B;
 71     F[val[rt]] = P + rt;
 72     for(int i = head[rt]; ~i; i = e[i].next){
 73         if(vis[e[i].to]) continue;
 74         calc(e[i].to,rt,1,1);
 75         calc(e[i].to,rt,val[rt],0);
 76     }
 77     for(int i = head[rt]; ~i; i = e[i].next)
 78         if(!vis[e[i].to]) solve(e[i].to);
 79 }
 80 int main() {
 81     init();
 82     int u,v;
 83     while(~scanf("%d%I64d",&n,&K)) {
 84         memset(head,-1,sizeof head);
 85         memset(vis,false,sizeof vis);
 86         tot = 0;
 87         for(int i = 1; i <= n; ++i)
 88             scanf("%I64d",val + i);
 89         for(int i = 1; i < n; ++i) {
 90             scanf("%d%d",&u,&v);
 91             add(u,v);
 92             add(v,u);
 93         }
 94         ans = INF;
 95         solve(1);
 96         if(ans == INF) puts("No solution");
 97         else printf("%d %d\n",ans.A,ans.B);
 98     }
 99     return 0;
100 }