SPOJ 913 [QTREE2]

gistfile1.cpp

#include <iostream>
#include <sstream>
#include <bitset>
#include <cstdio>
#include <string>
#include <cstring>
#include <vector>
#include <queue>
#include <stack>
#include <set>
#include <map>
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <cctype>
#include <numeric>
#include <list>
#define FOR(i,A) for(typeof (A).begin() i = (A).begin() ; i != (A).end() ; i++)
#define mp make_pair
#define debug( x ) cout << #x << " = " << x << endl
#define clr(v,x) memset( v, x , sizeof v )
#define all(x) (x).begin() , (x).end()
#define rall(x) (x).rbegin() , (x).rend()
#define f(i,a,b) for(int i = a ; i < b ; i++)
#define fd(i,a,b) for(int i = a ; i >= b ; i--)
#define PI acos( -1.0 )
#define EPS 1E-9
#define TAM 10010

using namespace std;

typedef long long ll ;
typedef long double ld ;
typedef pair<int,ll> ii ;
typedef pair<int,ii> pii ;

vector<ii> g[ TAM ] ;
ll P[ TAM ][ 25 ] ;
ll D[ TAM ][ 2 ] ;
ll L[ TAM ] ;
ll T[ TAM ] ;
ll n ;
ll lca ;

void dfs( ll x , ll lvl = 1 ){
	L[ x ] = lvl ;
	f( i , 0 , g[ x ].size() ){
		int u = g[ x ][ i ].first ;
		if( L[ u ] ) continue ;
		D[ u ][ 0 ] = D[ x ][ 0 ] + g[ x ][ i ].second ;
		D[ u ][ 1 ] = D[ x ][ 1 ] + 1 ;
		T[ u ] = x ;
		dfs( u , lvl + 1 ) ;
	}
}

void init(){
	f( i , 0 , n ){ L[ i ] = T[ i ] = D[ i ][ 0 ] = D[ i ][ 1 ] = 0 ; g[ i ].clear() ; }
	clr( P , 0 ) ;
}

void pre(){
	dfs( 0 ) ;
	f( i , 0 , n ) P[ i ][ 0 ] = T[ i ] ;
	for(int j = 1 ; (1<<j) <= n ; j++)
		f( i , 0 , n )
			P[ i ][ j ] = P[ P[ i ][ j - 1 ] ][ j - 1 ] ;
}

ll dist( ll x , ll y , ll dim = 0 ){
	if( x == y ){ lca = x ; return 0 ; }

	if( L[ x ] < L[ y ] ){ ll tmp = x ; x = y ; y = tmp ; }

	ll log = 0 ;
	for( log = 0 ; (1<<log) <= n ; log++ ) ;
	log-- ;

	ll R = 0 ;
	fd( i , log , 0 ){
		if( L[ x ] - (1<<i) >= L[ y ] ){
			R += ( D[ x ][ dim ] - D[ P[ x ][ i ] ][ dim ] ) ;
			x = P[ x ][ i ] ;
		}
	}

	if( x == y ){ lca = x ; return R ; }

	fd( i , log , 0 ){
		if( P[ x ][ i ] != P[ y ][ i ] ){
			R += ( D[ x ][ dim ] - D[ P[ x ][ i ] ][ dim ] ) ;
			R += ( D[ y ][ dim ] - D[ P[ y ][ i ] ][ dim ] ) ;
			x = P[ x ][ i ] , y = P[ y ][ i ] ;
		}
	}
	R += ( D[ x ][ dim ] - D[ P[ x ][ 0 ] ][ dim ] ) ;
	R += ( D[ y ][ dim ] - D[ P[ y ][ 0 ] ][ dim ] ) ;
	lca = P[ x ][ 0 ] ;
	return R ;
}

ll kth( ll x , ll y , ll k ){
	ll dt = dist( x , y , 1 ) ;
	ll LCA = lca ;
	ll dl = dist( x , LCA , 1 ) ;

	ll log = 0 ;
	for( log = 0 ; (1<<log) <= n ; log++ ) ;
	log-- ;

	ll resp ;
	k-- ;
	if( dl > k ){
		fd( i , log , 0 )
			if( L[ x ] - (1<<i) >= L[ LCA ] )
				if( k - (1<<i) >= 0 ){
					k -= (1<<i) ;
					x = P[ x ][ i ] ;
				}
		resp = x ;
	}else if( dl < k ){
		k = dt - k ;
		fd( i , log , 0 )
			if( L[ y ] - (1<<i) >= L[ LCA ] )
				if( k - (1<<i) >= 0 ){
					y = P[ y ][ i ] ;
					k -= (1<<i) ;
				}
		resp = y ;
	}else
		resp = LCA ;
	return resp + 1 ;
}

int main(){

	ll test , a , b , c ;
	char s[ 10 ] ;
	scanf("%lld" , &test ) ;
	while( test-- ){
		scanf("%lld" , &n ) ;
		init() ;
		f( i , 0 , n-1 ){
			scanf("%lld%lld%lld" , &a , &b , &c ) ;
			a-- ; b-- ;
			g[ a ].push_back( mp( b , c ) ) ;
			g[ b ].push_back( mp( a , c ) ) ;
		}
		pre() ;
		while( scanf("%s" , s ) == 1 ){
			if( s[ 1 ] == 'O' ) break ;
			ll R ;
			if( s[ 1 ] == 'I' ){
				scanf("%lld%lld" , &a , &b ) ;
				a-- ; b-- ;
				R = dist( a , b ) ;
			}else{
				scanf("%lld%lld%lld" , &a , &b , &c ) ;
				a-- ; b-- ;
				R = kth( a , b , c ) ;
			}
			printf("%lld\n" , R ) ;
		}
		printf("\n" ) ;
	}
    return 0 ;
}