SPOJ QTREE2 - Query on a tree II

Dist দিলে a,b এর মধ্যে Distance বের করতে হবে 

আর kth দিলে a থেকে b যেতে  kth তম নোড কোনটা?  

• DIST a b : ask for the distance between node a and node b
  or
• KTH a b k : ask for the k-th node on the path from node a to node b

1

6
1 2 1
2 4 1
2 5 2
1 3 1
3 6 2
DIST 4 6
KTH 4 6 4
DONE

Output:
5
3

#include <iostream>
#include <cstdio>
#include <cstring>
#define N 10005
using namespace std;
struct edge
{
    int to,next,w;
} e[N<<1];
int head[N],cnt;
void add(int u,int v,int w)
{
    e[++cnt]=(edge)
    {
        v,head[u],w
    };
    head[u]=cnt;
}
int n;
int deep[N],dis[N],p[N][30];
void dfs(int x,int fa)
{
    deep[x]=deep[fa]+1;
    p[x][0]=fa;
    for(int i=0; p[x][i]; i++)
        p[x][i+1]=p[p[x][i]][i];
    for(int i=head[x]; i; i=e[i].next)
    {
        int to=e[i].to;
        if(to==fa)
            continue;
        dis[to]=dis[x]+e[i].w;
        dfs(to,x);
    }
}
int lca(int a,int b)
{
    if(deep[a]<deep[b])
        swap(a,b);
    for(int j=14; j>=0; j--)
        if(deep[a]-(1<<j)>=deep[b])
            a=p[a][j];
    if(a==b)
        return a;
    for(int j=14; j>=0; j--)
        if(p[a][j]&&p[a][j]!=p[b][j])
        {
            a=p[a][j];
            b=p[b][j];
        }
    return p[a][0];
}
int get(int u,int d)
{
    for(int j=14; j>=0; j--)
        if(deep[u]-(1<<j)>=d)
            u=p[u][j];
    return u;
}
int main()
{
int t;
    cin>>t;
    while(t--)
    {
        memset(head,0,sizeof head);
        memset(p,0,sizeof p);
        cnt=0;
        cin>>n;
        for(int i=1; i<n; i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        dfs(1,0);
        char op[100];
        while(1)
        {
            scanf("%s",op);
            if(op[1]=='O')
                break;
            if(op[1]=='I')
            {
                int u,v;
                scanf("%d%d",&u,&v);
                printf("%d\n",dis[u]+dis[v]-2*dis[lca(u,v)]);
            }
            else
            {
                int u,v,k;
                scanf("%d%d%d",&u,&v,&k);
                int fa=lca(u,v),d;
                if(deep[u]-deep[fa]<k)
                {
                    d=deep[fa]+k-(deep[u]-deep[fa]+1);
                    u=v;
                }
                else
                    d=deep[u]-k+1;
                printf("%d\n",get(u,d));
            }
        }
    }
}