IOICJ 108 Where is He! ( Tree, Timestamp )

Problem Description:
You are given a tree with N nodes. There are Q queries, with u, v, w, where you want to know the node x such that x is a node in the simple path between u and v, where it is also the nearest to w.

Constraints:
1≤T≤50
2≤N,Q≤1e5
1≤a,b,u,v,w≤N

Sample Input:
1
6 2
1 2
2 4
4 5
2 3
3 6
1 6 5
4 5 1

Sample Output:
2
4

Solution:
If w is out of subtree of root lca( u, v ), then x is lca( u, v ).
Otherwise x is either lca( w, u ) or lca( w, v ).

#include <bits/stdc++.h>
using namespace std;

typedef vector< int > vi;
typedef vector< vi > vvi;

int T;

const int MAXN = ( int ) 1e5;

int N, Q;
vvi G;

void init(){
  cin >> N >> Q;
  G = vvi( N );
  for( int i = 0; i < N - 1; ++i ){
    int u, v; cin >> u >> v; --u, --v;
    G[ u ].emplace_back( v );
    G[ v ].emplace_back( u );
  }
}

int TS;
int tin[ MAXN ], tout[ MAXN ];
void ts_dfs( int u, int fa ){
  tin[ u ] = ++TS;
  for( int v : G[ u ] ){
    if( v == fa ) continue;
    ts_dfs( v, u );
  }
  tout[ u ] = TS;
}

int dpt[ MAXN ], par[ 20 ][ MAXN ];
void lca_dfs( int u, int fa, int d ){
  dpt[ u ] = d;
  par[ 0 ][ u ] = fa;
  for( int v : G[ u ] ){
    if( v == fa ) continue;
    lca_dfs( v, u, d + 1 );
  }
}

int lca( int u, int v ){
  if( not ( dpt[ u ] >= dpt[ v ] ) ) swap( u, v );
  int diff = dpt[ u ] - dpt[ v ];
  for( int i = 20 - 1; diff; --i ){
    if( diff >= 1 << i ){
      diff -= 1 << i;
      u = par[ i ][ u ];
    }
  }
  if( u == v ) return u;
  for( int i = 20 - 1; i >= 0; --i ){
    if( par[ i ][ u ] != par[ i ][ v ] ){
      u = par[ i ][ u ];
      v = par[ i ][ v ];
    }
  }
  return par[ 0 ][ u ];
}

void solve(){
  TS = 0;
  ts_dfs( 0, -1 );
  for( int i = 0; i < 20; ++i ){
    for( int j = 0; j < N; ++j ){
      par[ i ][ j ] = -1;
    }
  }
  lca_dfs( 0, -1, 0 );
  for( int i = 0; i + 1 < 20; ++i ){
    for( int j = 0; j < N; ++j ){
      int m = par[ i ][ j ];
      if( m == -1 ) continue;
      par[ i + 1 ][ j ] = par[ i ][ m ];
    }
  }
  for( int qi = 0; qi < Q; ++qi ){
    int u, v, w; cin >> u >> v >> w; --u, --v, --w;
    int f = lca( u, v );
    if( tin[ w ] < tin[ f ] or tout[ f ] < tin[ w ] ){
      cout << f + 1 << endl; // out of subtree
    } else{
      int a = lca( w, u );
      int b = lca( w, v );
      if( dpt[ w ] - dpt[ a ] < dpt[ w ] - dpt[ b ] ){
        cout << a + 1 << endl;
      } else{
        cout << b + 1 << endl;
      }
    }
  }
}

signed main(){
  ios::sync_with_stdio( 0 );
  cin >> T;
  for( int ti = 0; ti < T; ++ti ){
    init();
    solve();
  }
  return 0;
}