【SPOJ】913 Query on a tree II

  1 #include<cstdio>

  2 #include<cstring>

  3 #include<queue>

  4 #include<iostream>

  5 #define MAXN 10010

  6 #define MAXM 20010

  7 using namespace std;

  8 struct LCT {

  9     int bef[MAXN];

 10     int next[MAXN][2], pre[MAXN], key[MAXN], sum[MAXN], num[MAXN];

 11     void Init() {

 12         memset(next, 0, sizeof(next));

 13         memset(pre, 0, sizeof(pre));

 14     }

 15     inline void PushUp(int x) {

 16         sum[x] = sum[next[x][0]] + sum[next[x][1]] + key[x];

 17         num[x] = num[next[x][0]] + num[next[x][1]] + 1;

 18     }

 19     inline void Rotate(int x, int kind) {

 20         int y, z;

 21         y = pre[x];

 22         z = pre[y];

 23         next[y][!kind] = next[x][kind];

 24         pre[next[x][kind]] = y;

 25         next[z][next[z][1] == y] = x;

 26         pre[x] = z;

 27         next[x][kind] = y;

 28         pre[y] = x;

 29         PushUp(y);

 30     }

 31     void Splay(int x) {

 32         int rt;

 33         for (rt = x; pre[rt]; rt = pre[rt])

 34             ;

 35         if (x != rt) {

 36             bef[x] = bef[rt];

 37             bef[rt] = 0;

 38             while (pre[x]) {

 39                 if (next[pre[x]][0] == x)

 40                     Rotate(x, 1);

 41                 else

 42                     Rotate(x, 0);

 43             }

 44             PushUp(x);

 45         }

 46     }

 47     void Access(int x) {

 48         int father;

 49         for (father = 0; x; x = bef[x]) {

 50             Splay(x);

 51             pre[next[x][1]] = 0;

 52             bef[next[x][1]] = x;

 53             next[x][1] = father;

 54             pre[father] = x;

 55             bef[father] = 0;

 56             father = x;

 57             PushUp(x);

 58         }

 59     }

 60     int DIST(int x, int y) {

 61         Access(y);

 62         for (y = 0; x; x = bef[x]) {

 63             Splay(x);

 64             if (!bef[x])

 65                 return sum[y] + sum[next[x][1]];

 66             pre[next[x][1]] = 0;

 67             bef[next[x][1]] = x;

 68             next[x][1] = y;

 69             pre[y] = x;

 70             bef[y] = 0;

 71             y = x;

 72             PushUp(x);

 73         }

 74         return 0;

 75     }

 76     int Select(int x, int k) {

 77         while (num[next[x][0]] + 1 != k) {

 78             if (num[next[x][0]] + 1 > k)

 79                 x = next[x][0];

 80             else {

 81                 k -= num[next[x][0]] + 1;

 82                 x = next[x][1];

 83             }

 84         }

 85         return x;

 86     }

 87     int KTH(int y, int x, int k) {

 88         Access(y);

 89         for (y = 0; x; x = bef[x]) {

 90             Splay(x);

 91             if (!bef[x]) {

 92                 if (num[next[x][1]] + 1 == k)

 93                     return x;

 94                 else if (num[next[x][1]] + 1 > k)

 95                     return Select(next[x][1], num[next[x][1]] - k + 1);

 96                 return Select(y, k - num[next[x][1]] - 1);

 97             }

 98             pre[next[x][1]] = 0;

 99             bef[next[x][1]] = x;

100             next[x][1] = y;

101             pre[y] = x;

102             bef[y] = 0;

103             y = x;

104             PushUp(x);

105         }

106         return 0;

107     }

108 } lct;

109 int first[MAXN], next[MAXM], v[MAXM], cost[MAXM], e;

110 bool vis[MAXN];

111 int INT() {

112     int res;

113     char ch;

114     bool neg;

115     while (ch = getchar(), !isdigit(ch) && ch != '-')

116         ;

117     if (ch == '-') {

118         neg = true;

119         res = 0;

120     } else {

121         neg = false;

122         res = ch - '0';

123     }

124     while (ch = getchar(), isdigit(ch))

125         res = res * 10 + ch - '0';

126     return neg ? -res : res;

127 }

128 char CHAR() {

129     char ch, res;

130     while (ch = getchar(), !isalpha(ch))

131         ;

132     while (ch = getchar(), isalpha(ch))

133         res = ch;

134     return res;

135 }

136 void AddEdge(int x, int y, int val) {

137     v[e] = y;

138     cost[e] = val;

139     next[e] = first[x];

140     first[x] = e++;

141 }

142 void BFS(int x) {

143     int i, y;

144     queue<int> q;

145     memset(vis, false, sizeof(vis));

146     vis[x] = true;

147     q.push(x);

148     while (!q.empty()) {

149         x = q.front();

150         q.pop();

151         for (i = first[x]; i != -1; i = next[i]) {

152             y = v[i];

153             if (!vis[y]) {

154                 lct.bef[y] = x;

155                 lct.key[y] = cost[i];

156                 vis[y] = true;

157                 q.push(y);

158             }

159         }

160     }

161 }

162 int main() {

163     char ch;

164     int c, n, i, x, y, k;

165     c = INT();

166     while (c--) {

167         n = INT();

168         lct.Init();

169         memset(first, -1, sizeof(first));

170         for (e = 0, i = 1; i < n; i++) {

171             x = INT(), y = INT(), k = INT();

172             AddEdge(x, y, k);

173             AddEdge(y, x, k);

174         }

175         BFS(1);

176         while (ch = CHAR(), ch != 'E') {

177             x = INT(), y = INT();

178             if (ch == 'T')

179                 printf("%d\n", lct.DIST(x, y));

180             else {

181                 k = INT();

182                 printf("%d\n", lct.KTH(x, y, k));

183             }

184         }

185         putchar('\n');

186     }

187     return 0;

188 }