SPOJ 375 树链剖分 QTREE - Query on a tree

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;

const int maxn = 10000 + 10;

int n;
int tot;
vector<int> G[maxn];
int u[maxn], v[maxn], d[maxn];

int fa[maxn];
int top[maxn];
int id[maxn];
int L[maxn];
int son[maxn];
int sz[maxn];

void dfs(int u)
{
sz[u] = 1; son[u] = 0;
for(int i = 0; i < G[u].size(); i++)
{
int v = G[u][i];
if(v == fa[u]) continue;
fa[v] = u;
L[v] = L[u] + 1;
dfs(v);
sz[u] += sz[v];
if(sz[v] > sz[son[u]]) son[u] = v;
}
}

void dfs2(int u, int tp)
{
id[u] = ++tot;
top[u] = tp;
if(son[u]) dfs2(son[u], tp);
for(int i = 0; i < G[u].size(); i++)
{
int v = G[u][i];
if(v == fa[u] || v == son[u]) continue;
dfs2(v, v);
}
}

int maxv[maxn << 2];

void update(int o, int L, int R, int p, int val)
{
if(p < L || p > R) return ;
if(L == R) { maxv[o] = val; return ; }
int M = (L + R) / 2;
update(o<<1, L, M, p, val);
update(o<<1|1, M+1, R, p, val);
maxv[o] = max(maxv[o<<1], maxv[o<<1|1]);
}

int query(int o, int L, int R, int qL, int qR)
{
if(qR < L || qL > R) return 0;
if(qL <= L && R <= qR) return maxv[o];
int M = (L + R) / 2;
return max(query(o<<1, L, M, qL, qR), query(o<<1|1, M+1, R, qL, qR));
}

int QUERY(int u, int v)
{
int ans = 0;
int t1 = top[u], t2 = top[v];
while(t1 != t2)
{
if(L[t1] < L[t2]) { swap(u, v); swap(t1, t2); }
ans = max(ans, query(1, 1, tot, id[t1], id[u]));
u = fa[t1]; t1 = top[u];
}
if(u == v) return ans;
if(L[u] < L[v]) swap(u, v);
ans = max(ans, query(1, 1, tot, id[son[v]], id[u]));
return ans;
}

int main()
{
int T; scanf("%d", &T);
while(T--)
{
scanf("%d", &n);
for(int i = 1; i <= n; i++) G[i].clear();
fa[1] = L[1] = 0;
for(int i = 1; i < n; i++)
{
scanf("%d%d%d", u + i, v + i, d + i);
G[u[i]].push_back(v[i]);
G[v[i]].push_back(u[i]);
}

dfs(1);
tot = 0;
dfs2(1, 1);

memset(maxv, 0, sizeof(maxv));
for(int i = 1; i < n; i++)
{
if(L[u[i]] < L[v[i]]) swap(u[i], v[i]);
update(1, 1, tot, id[u[i]], d[i]);
}

char op[30];
while(scanf("%s", op) && op[0] != 'D')
{
int x, y; scanf("%d%d", &x, &y);
if(op[0] == 'Q')
{
printf("%d\n", QUERY(x, y));
}
else
{
update(1, 1, tot, id[u[x]], y);
}
}
}

return 0;
}