bzoj2588 Spoj 10628. Count on a tree

  1 #include<cstdio>
  2 #include<algorithm>
  3 #include<map>
  4 #define N 2000010
  5 using namespace std;
  6 map<int,int>ma;
  7 int n,m,i,a,b,c,e[N],last,d[N],id[N],deep[N],j; 
  8 int dp,pre[N],p[N],tt[N],l[N],r[N],ls[N],rs[N],s[N],tot,root[N];
  9 int jump[100100][19];
 10 void link(int x,int y)
 11 {
 12     dp++;pre[dp]=p[x];p[x]=dp;tt[dp]=y;
 13 }
 14 void build(int &x,int a,int b)
 15 {
 16     tot++;x=tot;
 17     l[x]=a;r[x]=b;
 18     if (b-a>1)
 19     {
 20         int m=(a+b)>>1;
 21         build(ls[x],a,m);
 22         build(rs[x],m,b);
 23     }
 24 }
 25 void insert(int &x,int y,int a,int b)
 26 {
 27     tot++;x=tot;
 28     l[x]=l[y];r[x]=r[y];
 29     ls[x]=ls[y];rs[x]=rs[y];s[x]=s[y];
 30     if ((a<=l[x])&&(r[x]<=b))
 31     {
 32         s[x]++;
 33         return;
 34     }
 35     int m=(l[x]+r[x])>>1;
 36     if (a<m) insert(ls[x],ls[y],a,b);
 37     if (m<b) insert(rs[x],rs[y],a,b);
 38     s[x]=s[ls[x]]+s[rs[x]];
 39 }
 40 void dfs(int x,int fa)
 41 {
 42     int i=p[x];
 43     insert(root[x],root[fa],e[x]-1,e[x]);
 44     jump[x][0]=fa;
 45     deep[x]=deep[fa]+1;
 46     while (i)
 47     {
 48         if (tt[i]!=fa) dfs(tt[i],x);
 49         i=pre[i];
 50     }
 51 }
 52 int query(int x,int y,int z,int w,int c)
 53 {
 54     if (r[x]-l[x]==1) 
 55     {
 56         return r[x];
 57     }
 58     int tmp=s[ls[x]]+s[ls[y]]-2*s[ls[z]];
 59     if ((l[ls[x]]<w)&&(w<=r[ls[x]])) tmp++;
 60     if (tmp>=c)
 61     return query(ls[x],ls[y],ls[z],w,c);
 62     else
 63     return query(rs[x],rs[y],rs[z],w,c-tmp);
 64 }
 65 int lca(int a,int b)
 66 {
 67     int i;
 68     if (deep[a]<deep[b]) a^=b^=a^=b;
 69     for (i=17;i>=0;i--)
 70     if (deep[jump[a][i]]>=deep[b]) a=jump[a][i];
 71     if (a==b) return a;
 72 
 73     for (i=17;i>=0;i--)
 74     if (jump[a][i]!=jump[b][i])
 75     {
 76         a=jump[a][i];
 77         b=jump[b][i];
 78     }
 79     return jump[a][0];
 80 }
 81 int main()
 82 {
 83     scanf("%d%d",&n,&m);
 84     for (i=1;i<=n;i++)
 85     {
 86         scanf("%d",&e[i]);
 87         d[i]=e[i];
 88     }
 89     sort(d+1,d+1+n);int cnt=0;d[0]=-0x37373737;
 90     for (i=1;i<=n;i++)
 91     {
 92         if (d[i]!=d[i-1]) cnt++;
 93         ma[d[i]]=cnt;
 94         id[cnt]=d[i];
 95     }
 96     for (i=1;i<=n;i++) e[i]=ma[e[i]];
 97     for (i=1;i<n;i++)
 98     {
 99         scanf("%d%d",&a,&b);
100         link(a,b);
101         link(b,a);
102     }
103     build(root[0],0,cnt);
104     dfs(1,0);
105 
106     for (i=1;i<=17;i++)
107     for (j=1;j<=n;j++)
108     jump[j][i]=jump[jump[j][i-1]][i-1];
109     
110     
111     int ans=0;
112     for (i=1;i<=m;i++)
113     {
114         scanf("%d%d%d",&a,&b,&c);
115         a^=ans;
116         int LCA=lca(a,b);
117         //printf("%d %d %d\n",s[root[a]],s[root[b]],LCA);
118         ans=id[query(root[a],root[b],root[LCA],e[LCA],c)];
119         printf("%d",ans);
120         if (i!=m) printf("\n");
121     }
122     
123 }