SPOJ 375 Query on a tree(树链剖分)

题目分析

本题主要是修改树上线的cost值，还有询问2点之间的最大的cost.因为本题数据量大，并且修改和查询次数多，因此我们需要使用线段树进行操作。因此我们需要对树进行剖分，并将每个点对应的线段进行更新即可，根节点没有对应线段，因为n个节点的树只有n-1条边，因此统一将线段与连接2端的深度较大的点对应。至于树链剖分的理解有一个很不错的博客博客地址，多看看配合在草纸上画出对应图形就能正确理解了。

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
#define mid (L+R)/2
#define lson o<<1, L, mid
#define rson o<<1|1, mid+1, R
const int maxn = 10005;
int ans[maxn<<2],edge[maxn][3];

void update(int o,int L,int R,int p,int val){  //线段树单点更新
    if(L == R){
        ans[o] = val;
        return ;
    }
    if(p <= mid) update(lson, p, val);
    else update(rson, p, val);
    ans[o] = max(ans[o<<1], ans[o<<1|1]);
}

int query(int o,int L,int R,int l,int r){  //线段树查询区间最大值
    if(l <= L && R <= r) return ans[o];
    int ret = 0;
    if(l <= mid) ret = max(ret, query(lson, l, r));
    if(r > mid) ret = max(ret, query(rson, l, r));
    return ret;
}

struct Edge{  //邻接表
    int to,next;
}e[maxn<<1];

int siz[maxn],dep[maxn],fa[maxn],son[maxn],top[maxn],w[maxn],head[maxn],tot,cnt;
//以该点为根节点的子树的节点数， 节点的深度， 节点的爸爸， 节点的在重链中的儿子，节点重链中的顶点，节点对应的边在线段树中的位置。

void addedge(int from,int to){ //邻接表中加边
    e[cnt].to = to; e[cnt].next = head[from]; head[from] = cnt++;
}

void dfs1(int now){ //计算son,siz,fa,dep数组
    son[now] = 0;
    siz[now] = 1;
    for(int u = head[now]; u != -1; u = e[u].next){
        int v = e[u].to;
        if(v != fa[now]){
            fa[v] = now;
            dep[v] = dep[now]+1;
            dfs1(v);
            siz[now] += siz[v];
            if(siz[v] > siz[son[now]])son[now] = v;
        }
    }
}

void dfs2(int now,int tp){ //计算top，w
    w[now] = ++tot;
    top[now] = tp;
    if(son[now]) dfs2(son[now], top[now]);
    for(int u = head[now]; u != -1; u = e[u].next){
        int v = e[u].to;
        if(v != son[now] && v != fa[now])
            dfs2(v, v);
    }
}

void init(){
    tot = cnt = fa[1] = dep[1] = 0;
    memset(head, -1, sizeof(head));
    memset(ans, 0, sizeof(ans));
    memset(siz, 0, sizeof(siz));
}

int solve(int a,int b){
    int f1 = top[a], f2 = top[b], ret = 0;
    while(f1 != f2){ 
        if(dep[f1] < dep[f2]){ //找到深度较大的重链的顶点，也有可能是轻链，是轻链是这时候top值是其本身
            swap(f1, f2);
            swap(a, b);
        }
        ret = max(ret, query(1, 1, tot, w[f1], w[a]));
        a = fa[f1];
        f1 = top[a];
    }
    if(a == b) return ret;
    if(dep[a] > dep[b]) swap(a,b); //在同一条重链上
    ret = max(ret, query(1, 1, tot, w[son[a]], w[b]));
    return ret;
}

int main(){
    int T,n;
    char cmd[10];
    scanf("%d", &T);
    while(T--){
        init();
        scanf("%d", &n);
        for(int i = 1; i < n; i++){
            scanf("%d%d%d", &edge[i][0], &edge[i][1], &edge[i][2]);
            addedge(edge[i][0], edge[i][1]);
            addedge(edge[i][1], edge[i][0]);
        }
        dfs1(1);
        dfs2(1, 1);
        for(int i = 1; i < n; i++){
            if(dep[edge[i][0]] > dep[edge[i][1]]) swap(edge[i][0], edge[i][1]);  //因为深度大的在线段树中的位置在一条链中深度小的后面
            update(1, 1, tot, w[edge[i][1]], edge[i][2]);
        }
        while(scanf("%s", cmd) != EOF && cmd[0] != 'D'){
            int from, to;
            if(cmd[0] == 'Q'){
                scanf("%d%d", &from, &to);
                printf("%d\n", solve(from, to));
            }
            else{
                scanf("%d%d", &from, &to);
                update(1, 1, tot, w[edge[from][1]], to);
            }
        }
    }
    return 0;
}