spoj 375--Query On a Tree [树链剖分]

#include <cstdio>
#include <algorithm>
using namespace std;
// 读入优化
inline int getInt()
{
    int res = 0;
    char ch;
    for (ch = '#'; ch < '0' || ch > '9'; ch = getchar());
    for (; ch >= '0' && ch <= '9'; ch = getchar()) 
        res = res * 10 + (int) ch - (int) '0';
    return res;
}
// 得到题目的操作
inline char getCommend()
{
    char ch, t;
    for (ch = ' '; ch == ' ' || ch == '\n'; ch = getchar());
    t = ch;
    for (; ch != ' ' && ch != '\n'; ch = getchar());
    return t;
}
const int inf = 1000007;        // 无穷大
const int N = 10007;            // 最多的结点个数
int n, edgeval
;                // 每条边的初始值
// 数组模拟指针，from为头，to为连接的点，next为下一个。nedge为边数
int from
, to[N << 1], next[N << 1], nedge;
// 每条边依附的点
int edgepos
;
// 插入边
void Insert(int a, int b)
{
    to[nedge] = b;
    next[nedge] = from[a];
    from[a] = nedge ++;
}
// 读入
void Init()
{
    memset(from, -1, sizeof(from));
    nedge = 0;
    n = getInt();
    for (int i = 0; i + 1 < n; i ++) 
    {
        int a = getInt() - 1, b = getInt() - 1;
        edgeval[i] = getInt();
        Insert(a, b);
        Insert(b, a);
    }
}
int top
, len
, npath;
int belong
, idx
, dep
, size
, father
, Q
;
void Split()
{
    npath = 0;
    int lo = 0, hi = 0;
    Q[0] = 0;
    dep[0] = 0;
    father[0] = -1;
    
    while (lo <= hi)
    {
        int u = Q[lo ++];
        for (int e = from[u]; e != -1; e = next[e])
        {
            int v = to[e];
            if (v != father[u])
            {
                edgepos[e >> 1] = v;
                Q[++ hi] = v;
                dep[v] = dep[u] + 1;
                father[v] = u;
            }
        }
    }
    
    for (int i = n - 1; i >= 0; i --)
    {
        // p为重儿子的编号
        int u = Q[i], p = -1;
        size[u] = 1;
        for (int e = from[u]; e != -1; e = next[e])
        {
            int v = to[e];
            if (v != father[u])
            {
                size[u] += size[v];
                if (p == -1 || size[v] > size[p]) // 求重儿子
                    p = v;
            }
        }
        
        if (p == -1)
        {
            belong[u] = npath;
            len[npath] = 1;
            idx[u] = 0;
            top[npath ++] = u;
        }
        else 
        {
            int x = belong[p];
            belong[u] = x;
            idx[u] = len[x] ++;
            top[x] = u;
        }
    }
}
int nnode;
// 线段树的结点
struct Node
{
    Node *lch, *rch;
    int lo, hi, dat;
    inline int mi()
    {
        return (lo + hi) >> 1;
    }
}node[N << 1], *tree
;
void Build(Node *p, int lo, int hi)
{
    p -> lo = lo;
    p -> hi = hi;
    p -> dat = - inf;
    if (lo + 1 < hi)
    {
        int mi = p -> mi();    
        p -> lch = &node[nnode ++];
        p -> rch = &node[nnode ++];
        Build(p -> lch, lo, mi);
        Build(p -> rch, mi, hi);
    }
}
void Modify(Node *p, int pos, int val)
{
    if (p -> lo + 1 == p -> hi) 
        p -> dat = val;
    else 
    {
        int mi = p -> mi();
        if (pos < mi) Modify(p -> lch, pos, val);
        else Modify(p -> rch, pos, val);
        p -> dat = max(p -> lch -> dat, p -> rch -> dat);
    }
}
int Ask(Node *p, int le, int ri)
{
    if (le <= p -> lo && ri >= p -> hi) 
        return p -> dat;
    else 
    {
        int mi = p -> mi();
        int ret = - inf;
        if (le < mi) ret = max(ret, Ask(p -> lch, le, ri));
        if (ri > mi) ret = max(ret, Ask(p -> rch, le, ri));
        return ret;
    }
}
void Prepare()
{
    nnode = 0;
    for (int i = 0; i < npath; i ++)
    {
        tree[i] = &node[nnode ++];
        Build(tree[i], 0, len[i]);
    }
    // 初始化边的权值
    for (int i = 0; i + 1 < n; i ++)
    {
        int u = edgepos[i];
        Modify(tree[belong[u]], idx[u], edgeval[i]);
    }
}
// 对于两点操作
int Find(int a, int b)
{
    int res = - inf;
    // x是点a所在的重链，y是点b所在的重链
    int x = belong[a], y = belong[b];
    
    while (x != y)
    {
        // 选择顶端元素深度较深的往上跳
        if (dep[top[x]] < dep[top[y]])
        {
            swap(a, b);
            swap(x, y);
        }
        res = max(res, Ask(tree[x], idx[a], len[x]));
        // 跳的过程
        a = father[top[x]];
        x = belong[a];
    }
    // 直到在同一重链
    if (idx[a] != idx[b])
    {
        if (idx[a] > idx[b]) swap(a, b);
        res = max(res, Ask(tree[x], idx[a], idx[b]));
    }
    return res;
}
void Solve()
{
    while (true)
    {
        char commend = getCommend();
        switch (commend) 
        {
            case 'C':
                int e = getInt() - 1, v = getInt();
                int u = edgepos[e];
                Modify(tree[belong[u]], idx[u], v);
                break;
            case 'Q':
                int a = getInt() - 1, b = getInt() - 1;
                printf("%d\n", Find(a, b));
                break;
            case 'D':
                return;
                break;
        }
    }
}
int main()
{
    for (int T = getInt(); T; T --)
    {
        Init();
        Split();
        Prepare();
        Solve();
    }
    
    return 0;
}