SPOJ QTREE2 Query on a tree II

  1 //minamoto
  2 #include<bits/stdc++.h>
  3 using namespace std;
  4 #define getc() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++)
  5 char buf[1<<21],*p1=buf,*p2=buf;
  6 inline int read(){
  7     #define num ch-‘0‘
  8     char ch;bool flag=0;int res;
  9     while(!isdigit(ch=getc()))
 10     (ch==‘-‘)&&(flag=true);
 11     for(res=num;isdigit(ch=getc());res=res*10+num);
 12     (flag)&&(res=-res);
 13     #undef num
 14     return res;
 15 }
 16 inline bool isop(char ch){
 17     return ch==‘I‘||ch==‘H‘||ch==‘O‘;
 18 }
 19 inline char readop(){
 20     char ch;
 21     while(!isop(ch=getc()));
 22     return ch;
 23 }
 24 const int N=10005;
 25 int head[N],Next[N<<1],edge[N<<1],ver[N<<1],fa[N][16],d[N],dist[N];
 26 int n,m,tot;
 27 inline void add(int u,int v,int e){
 28     ver[++tot]=v,Next[tot]=head[u],head[u]=tot,edge[tot]=e;
 29     ver[++tot]=u,Next[tot]=head[v],head[v]=tot,edge[tot]=e;
 30 }
 31 void dfs(int u,int f){
 32     d[u]=d[f]+1,fa[u][0]=f;
 33     for(int i=1;(1<<i)<=d[u];++i) fa[u][i]=fa[fa[u][i-1]][i-1];
 34     for(int i=head[u];i;i=Next[i]){
 35         int v=ver[i];
 36         if(v==f) continue;
 37         dist[v]=dist[u]+edge[i],dfs(v,u);
 38     }
 39 }
 40 int LCA(int x,int y){
 41     if(d[x]<d[y]) swap(x,y);
 42     for(int i=15;i>=0;--i)
 43     if(d[fa[x][i]]>=d[y]) x=fa[x][i];
 44     if(x==y) return x;
 45     for(int i=15;i>=0;--i)
 46     if(fa[x][i]!=fa[y][i])
 47     x=fa[x][i],y=fa[y][i];
 48     return fa[x][0];
 49 }
 50 int querylen(int x,int y,int k){
 51     int lca=LCA(x,y);
 52     if(d[x]-d[lca]+1>=k){
 53         int ans=d[x]-k+1;
 54         for(int i=15;i>=0;--i)
 55         if((1<<i)<=d[x]-ans) x=fa[x][i];
 56         return x;
 57     }
 58     else{
 59         int ans=d[lca]*2+k-d[x]-1;
 60         for(int i=15;i>=0;--i){
 61             if(d[y]-(1<<i)>=ans) y=fa[y][i];
 62         }
 63         return y;
 64     }
 65 }
 66 int main(){
 67     //freopen("testdata.in","r",stdin);
 68     int q=read();
 69     while(q--){
 70         memset(head,0,sizeof(head));
 71         memset(fa,0,sizeof(fa));
 72         d[1]=dist[1]=0,tot=0;
 73         int n=read();
 74         for(int i=1;i<n;++i){
 75             int u=read(),v=read(),e=read();
 76             add(u,v,e);
 77         }
 78         dfs(1,0);
 79         bool flag=true;
 80         while(flag){
 81             char op=readop();
 82             switch(op){
 83                 case ‘O‘:flag=false;break;
 84                 case ‘I‘:{
 85                     int u=read(),v=read();
 86                     int lca=LCA(u,v);
 87                     printf("%d\n",dist[u]+dist[v]-2*dist[lca]);
 88                     break;
 89                 }
 90                 case ‘H‘:{
 91                     int u=read(),v=read(),k=read();
 92                     printf("%d\n",querylen(u,v,k));
 93                     break;
 94                 }
 95             }
 96         }
 97         puts("");
 98     }
 99     return 0;
100 }