SPOJ QTREE2 Query on a tree II

Description

You are given a tree (an undirected acyclic connected graph) with N
nodes, and edges numbered 1, 2, 3…N-1. Each edge has an integer
value assigned to it, representing its length.

We will ask you to perfrom some instructions of the following form:

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

Example: N = 6 1 2 1 // edge connects node 1 and node 2 has cost 1 2 4
1 2 5 2 1 3 1 3 6 2

Path from node 4 to node 6 is 4 -> 2 -> 1 -> 3 -> 6 DIST 4 6 : answer
is 5 (1 + 1 + 1 + 2 = 5) KTH 4 6 4 : answer is 3 (the 4-th node on the
path from node 4 to node 6 is 3) Input

The first line of input contains an integer t, the number of test
cases (t <= 25). t test cases follow.

For each test case:

In the first line there is an integer N (N <= 10000)
In the next N-1 lines, the i-th line describes the i-th edge: a line with three integers a b c denotes an edge between a, b of cost c

(c <= 100000)
The next lines contain instructions “DIST a b” or “KTH a b k”
The end of each test case is signified by the string “DONE”.

There is one blank line between successive tests. Output

For each “DIST” or “KTH” operation, write one integer representing its
result.

Print one blank line after each test.

倍增处理出祖先和距离，距离只要在求lca的时候顺便累加，第k个点可以判断在lca之前还是lca之后，然后从第k个点所在的那一侧倍增。

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define M(a) memset(a,0,sizeof(a))
const int mx=18;
char s[100];
int fir[100010],ne[200010],to[200010],len[200010],anc[100010][20],dis[100010][20],dep[100010],n;
void add(int num,int f,int t,int l)
{
    ne[num]=fir[f];
    fir[f]=num;
    to[num]=t;
    len[num]=l;
}
void init()
{
    int i,x,y,z;
    M(fir);
    M(ne);
    M(anc);
    M(dis);
    scanf("%d",&n);
    for (i=1;i<n;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        add(i*2,x,y,z);
        add(i*2+1,y,x,z);
    }
}
void dfs(int u,int fa,int d)
{
    int i,j,k,v,x,y,z;
    dep[u]=d;
    for (i=fir[u];i;i=ne[i])
      if ((v=to[i])!=fa)
      {
        anc[v][0]=u;
        dis[v][0]=len[i];
        dfs(v,u,d+1);
      }
}
void make()
{
    int i,k;
    for (k=1;k<=mx;k++)
      for (i=1;i<=n;i++)
      {
        anc[i][k]=anc[anc[i][k-1]][k-1];
        dis[i][k]=dis[i][k-1]+dis[anc[i][k-1]][k-1];
      }
}
int q_dis()
{
    int i,k,u,v,ans=0;
    scanf("%d%d",&u,&v);
    if (dep[u]<dep[v]) swap(u,v);
    if (dep[u]>dep[v])
    {
        for (k=mx;k>=0;k--)
          if (dep[v]+(1<<k)<=dep[u])
          {
            ans+=dis[u][k];
            u=anc[u][k];
          }
    }
    if (u==v) return ans;
    for (k=mx;k>=0;k--)
      if (anc[u][k]!=anc[v][k])
      {
        ans+=dis[u][k]+dis[v][k];
        u=anc[u][k];
        v=anc[v][k];
      }
    return ans+dis[u][0]+dis[v][0];
}
int lca(int u,int v)
{
    int k;
    if (dep[u]<dep[v]) swap(u,v);
    if (dep[u]>dep[v])
    {
        for (k=mx;k>=0;k--)
          if (dep[v]+(1<<k)<=dep[u])
            u=anc[u][k];
    }
    if (u==v) return u;
    for (k=mx;k>=0;k--)
      if (anc[u][k]!=anc[v][k])
      {
        u=anc[u][k];
        v=anc[v][k];
      }
    return anc[u][0];
}
int q_anc(int u,int p)
{
    int k;
    for (k=mx;k>=0;k--)
      if ((1<<k)&p)
        u=anc[u][k];
    return u;
}
int q_v()
{
    int i,j,k,u,v,x,y,z,ans,rk;
    scanf("%d%d%d",&u,&v,&rk);
    rk--;
    x=lca(u,v);
    if (dep[u]-dep[x]>=rk) return q_anc(u,rk);
    return q_anc(v,dep[v]+dep[u]-rk-2*dep[x]);
}
int main()
{
    int T;
    scanf("%d",&T);
    while (T--)
    {
        init();
        dfs(1,0,1);
        make();
        while (scanf("%s",s)&&s[1]!='O')
          if (s[1]=='I')
            printf("%d\n",q_dis());
          else printf("%d\n",q_v());
    }
}