SPOJ QTREE2 Query on a tree II
Query on a tree II
64-bit integer IO format: %lld Java class name: Main
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.
Example
Input: 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
Source
1 #include <bits/stdc++.h> 2 using namespace std; 3 const int maxn = 100010; 4 struct arc { 5 int to,w,next; 6 arc(int x = 0,int y = 0,int z = -1) { 7 to = x; 8 w = y; 9 next = z; 10 } 11 } e[maxn<<1]; 12 int head[maxn],tot; 13 void add(int u,int v,int w) { 14 e[tot] = arc(v,w,head[u]); 15 head[u] = tot++; 16 } 17 struct LCT { 18 int fa[maxn],ch[maxn][2],sz[maxn]; 19 int key[maxn],pa[maxn],sum[maxn]; 20 inline void pushup(int x) { 21 sz[x] = 1 + sz[ch[x][0]] + sz[ch[x][1]]; 22 sum[x] = key[x] + sum[ch[x][0]] + sum[ch[x][1]]; 23 } 24 void rotate(int x,int kd) { 25 int y = fa[x]; 26 ch[y][kd^1] = ch[x][kd]; 27 fa[ch[x][kd]] = x; 28 fa[x] = fa[y]; 29 ch[x][kd] = y; 30 fa[y] = x; 31 if(fa[x]) ch[fa[x]][y == ch[fa[x]][1]] = x; 32 pushup(y); 33 } 34 void splay(int x,int goal) { 35 while(fa[x] != goal) { 36 if(fa[fa[x]] == goal) rotate(x,x == ch[fa[x]][0]); 37 else { 38 int y = fa[x],z = fa[y],s = (y == ch[z][0]); 39 if(x == ch[y][s]) { 40 rotate(x,s^1); 41 rotate(x,s); 42 } else { 43 rotate(y,s); 44 rotate(x,s); 45 } 46 } 47 } 48 pushup(x); 49 } 50 void access(int x) { 51 for(int y = 0; x; x = pa[x]) { 52 splay(x,0); 53 fa[ch[x][1]] = 0; 54 pa[ch[x][1]] = x; 55 ch[x][1] = y; 56 fa[y] = x; 57 pa[y] = 0; 58 y = x; 59 pushup(x); 60 } 61 } 62 int select(int x,int k) { 63 while(sz[ch[x][0]] + 1 != k) { 64 if(k < sz[ch[x][0]] + 1) x = ch[x][0]; 65 else { 66 k -= sz[ch[x][0]] + 1; 67 x = ch[x][1]; 68 } 69 } 70 return x; 71 } 72 int kth(int y,int x,int k) { 73 access(y); 74 for(y = 0; x; x = pa[x]) { 75 splay(x,0); 76 if(!pa[x]) { 77 if(sz[ch[x][1]] + 1 == k) return x; 78 if(sz[ch[x][1]] + 1 > k) return select(ch[x][1],sz[ch[x][1]] - k + 1); 79 return select(y,k - sz[ch[x][1]] - 1); 80 } 81 fa[ch[x][1]] = 0; 82 pa[ch[x][1]] = x; 83 ch[x][1] = y; 84 fa[y] = x; 85 pa[y] = 0; 86 y = x; 87 pushup(x); 88 } 89 return 0; 90 } 91 int dis(int x,int y) { 92 access(y); 93 for(y = 0; x; x = pa[x]) { 94 splay(x,0); 95 if(!pa[x]) return sum[y] + sum[ch[x][1]]; 96 fa[ch[x][1]] = 0; 97 pa[ch[x][1]] = x; 98 ch[x][1] = y; 99 fa[y] = x; 100 pa[y] = 0; 101 y = x; 102 pushup(x); 103 } 104 return 0; 105 } 106 void build(int _val,int u,int v) { 107 sz[v] = 1; 108 ch[v][1] = ch[v][0] = fa[v] = 0; 109 pa[v] = u; 110 key[v] = sum[v] = _val; 111 } 112 } spt; 113 void dfs(int u,int fa) { 114 for(int i = head[u]; ~i; i = e[i].next) { 115 if(e[i].to == fa) continue; 116 spt.build(e[i].w,u,e[i].to); 117 dfs(e[i].to,u); 118 } 119 } 120 int main() { 121 int kase,n,u,v,w; 122 char op[10]; 123 scanf("%d",&kase); 124 while(kase--) { 125 scanf("%d",&n); 126 memset(head,-1,sizeof head); 127 tot = 0; 128 for(int i = 1; i < n; ++i) { 129 scanf("%d%d%d",&u,&v,&w); 130 add(u,v,w); 131 add(v,u,w); 132 } 133 spt.build(0,0,1); 134 dfs(1,1); 135 while(scanf("%s",op) && op[1] != 'O') { 136 if(op[1] == 'I') { 137 scanf("%d%d",&u,&v); 138 if(u == v) puts("0"); 139 else printf("%d\n",spt.dis(u,v)); 140 } else if(op[1] == 'T') { 141 scanf("%d%d%d",&u,&v,&w); 142 printf("%d\n",spt.kth(u,v,w)); 143 } else break; 144 } 145 puts(""); 146 } 147 return 0; 148 } 149 /* 150 151 1 152 5 153 1 2 1 154 2 3 2 155 3 4 3 156 3 5 4 157 158 */