913. Query on a tree II

http://www.spoj.com/problems/QTREE2/

コスト付き木が与えられる。
頂点a,b間のコストを答えるクエリと
頂点a,b間のパスのk番目の頂点を答えるクエリを処理しろというような問題。

Heavy-Light Decompositionを用いて、LCAなどを処理するようにしました。
速度はそこまで早くないようです。
もう少しコンパクトな実装にしたいです。

struct Heavylight{
	vector<vector<PI> > G;
	vector<int> sizevis;
	vector<vector<int> > segments;
	vector<PI> pos;
	vector<int> parent;
	vector<int> segdist;

	void sizedfs(int v){
		queue<int> q;
		q.push(v);
		vector<int> near;
		sizevis = vector<int>(G.size());
		while(!q.empty()){
			int cv = q.front();
			q.pop();
			if(sizevis[cv]) continue;
			sizevis[cv] = 1;
			near.pb(cv);
			FOR(it,G[cv]){
				int to=it->F;
				if(sizevis[to]) continue;
				q.push(to);
			}
		}
		
		sizevis = vector<int>(G.size());
		
		FORR(it,near){
			int ret=1;
			FOR(ch,G[*it]){
				int to=ch->F;
				if(sizevis[to]) ret += sizevis[to];
				else parent[*it] = to;
			}
			sizevis[*it] = ret;
		}
	}
	
	void makesegment(int v){
		vector<int> childsegment;
		segments.push_back(vector<int>(1,v));
		vector<int>& currentsegment = segments.back();
		while(true){
			vector<PI> size;
			FOR(it,G[v])
				if(sizevis[it->F]<sizevis[v])
					size.pb(mp(-sizevis[it->F],it-G[v].begin()));
			sort(ALL(size));
			for(int i=1;i<SZ(size);++i){
				const PI& e=G[v][size[i].S];
				childsegment.pb(e.F);
				segdist[e.F] = e.S;
			}
			
			if(SZ(size)){
				const PI& e = G[v][size[0].S];
				int chv = e.F;
				currentsegment.pb(chv);
				segdist[chv] = segdist[v] + e.S;
				v = chv;
			}else break;
		}
		
		for(int i=0;i<SZ(currentsegment);++i)
			pos[currentsegment[i]] = mp(segments.size()-1,i);
		
		FOR(it,childsegment)
			makesegment(*it);
	}

	void init(int root){
		parent = vector<int>(G.size(),-1);
		pos = vector<PI>(G.size());
		segdist = vector<int>(G.size());
		sizedfs(root);
		makesegment(root);
	}
	
	Heavylight(const vector<vector<PI> >& aG,int root=0):G(aG){
		init(root);
	}

	int lca(int u,int v){
		while(true){
			if(pos[u].F == pos[v].F) return pos[u].S>pos[v].S?v:u;
			if(sizevis[segments[pos[u].F][0]] >
				 sizevis[segments[pos[v].F][0]]) swap(u,v);
			u = parent[segments[pos[u].F][0]];
		}
	}

	int dist(int u,int v){
		int ret=0;
		while(true){
			if(pos[u].F == pos[v].F) return ret + abs(segdist[u]-segdist[v]);
			if(sizevis[segments[pos[u].F][0]] >
				 sizevis[segments[pos[v].F][0]])
				swap(u,v);
			ret += segdist[u];
			u = parent[segments[pos[u].F][0]];
		}
	}

	int kth(int u,int v,int k){// u=v0,v=vn,return vk
		int lcav = lca(u,v);
		while(pos[u].F != pos[lcav].F){
			if(pos[u].S>=k) return segments[pos[u].F][pos[u].S-k];
			k -= pos[u].S+1;
			u = parent[segments[pos[u].F][0]];

		}
    
		if(pos[u].S-pos[lcav].S>=k)
			return segments[pos[u].F][pos[u].S-k];

		k -= pos[u].S - pos[lcav].S;
    
		int tv = v;
		int di = 0;

		while(pos[tv].F != pos[lcav].F){
			di += pos[tv].S+1;
			tv = parent[segments[pos[tv].F][0]];
		}
    
		di += pos[tv].S - pos[lcav].S;
		k = di - k;
		while(pos[v].F != pos[lcav].F){
			if(pos[v].S>=k) return segments[pos[v].F][pos[v].S-k];
			k -= pos[v].S+1;
			v = parent[segments[pos[v].F][0]];
		}
    
		if(pos[v].S-pos[lcav].S>=k)
			return segments[pos[v].F][pos[v].S-k];
	}
};


inline int in(){
	int n;
	scanf("%d",&n);
	return n;
}


void solve(){
	int n=in();
	vector<vector<PI> > G(n);
  
	rep(i,n-1){
		int a,b,c;
		a=in()-1;
		b=in()-1;
		c=in();
		G[a].pb(mp(b,c));
		G[b].pb(mp(a,c));
	}

	Heavylight hl(G);
	char q[10];
	while(true){
		int a,b;
		scanf(" %s",q);
		if(q[1]=='I'){
			a=in();
			b=in();
			printf("%d\n",hl.dist(a-1,b-1));
		}else if(q[1]=='T'){
			int k;
			a=in();
			b=in();
			k=in();
			printf("%d\n",hl.kth(a-1,b-1,k-1) + 1);
		}else break;
	}
}

int main(int argc, char *argv[])
{
	int t=in();
	rep(i,t) solve();
	return 0;
}