bzoj1036: [ZJOI2008]树的统计Count - 树链剖分

题目链接： http://www.lydsy.com/JudgeOnline/problem.php?id=1036

题意： 

给一棵树， 点上有权值， 带单点更新， 询问两点间权值最大值以及路径上所有点权值和。

树链剖分入门题。

#include <bits/stdc++.h>
using namespace std;
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1

typedef long long ll;
const int maxn = 30000 + 10;
const int inf = 1 << 30;

struct Edge{
    int to, next;
}e[maxn*2];

int n;
int cnt, head[maxn], w[maxn];

//
int fa[maxn], size[maxn], dep[maxn];
int top[maxn];
int son[maxn];
int idx[maxn];
int pos;

//
int Max[maxn<<2], Sum[maxn<<2];

void init()
{
    cnt = pos = 0;
    memset(head, -1, sizeof(head));
    memset(son, -1, sizeof(son));
}

void add(int u, int v)
{
    e[cnt].to = v;
    e[cnt].next = head[u];
    head[u] = cnt++;
}

void dfs1(int u, int pre, int deep)
{
    fa[u] = pre; size[u] = 1; dep[u] = deep;
    for(int i=head[u]; i!=-1; i=e[i].next)
    {
        int v = e[i].to;
        if(v == pre) continue;
        dfs1(v, u, deep+1);
        size[u] += size[v];
        if(son[u] == -1 || size[v] > size[son[u]]) son[u] = v;
    }
}

void dfs2(int u, int anc)
{
    top[u] = anc;
    if(son[u] != -1)
    {
        idx[u] = pos++;
        dfs2(son[u], anc);
    }
    else
    {
        idx[u] = pos++;
        return;
    }
    for(int i=head[u]; i!=-1; i=e[i].next)
    {
        int v = e[i].to;
        if(v != fa[u] && v!= son[u]) dfs2(v, v);
    }
}

void build(int l, int r, int rt)
{
    Max[rt] = -inf;
    Sum[rt] = 0;
    if(l == r) return;
    int m = (l+r) >> 1;
    build(lson);
    build(rson);
}

void pushUp(int rt)
{
    Max[rt] = max(Max[rt<<1], Max[rt<<1|1]);
    Sum[rt] = Sum[rt<<1] + Sum[rt<<1|1];
}

void update(int p, int v, int l, int r, int rt)
{
    if(l == r)
    {
        Max[rt] = v;
        Sum[rt] = v;
        return;
    }
    int m = (l+r) >> 1;
    if( p <= m) update(p, v, lson);
    else update(p, v, rson);
    pushUp(rt);
}

int query_sum(int L, int R, int l, int r, int rt)
{
    if( L<= l && r <= R) return Sum[rt];
    int ret = 0;
    int m = (l + r) >> 1;
    if(L <= m) ret += query_sum(L, R, lson);
    if(R > m) ret += query_sum(L, R, rson);
    return ret;
}

int query_max(int L, int R, int l, int r, int rt)
{
    if( L<= l && r <= R) return Max[rt];
    int ret = -inf;
    int m = (l + r) >> 1;
    if(L <= m) ret = max(ret, query_max(L, R, lson) );
    if(R > m) ret = max(ret, query_max(L, R, rson) );
    return ret;
}

int find_sum(int u, int v)
{
    int f1 = top[u], f2 = top[v];
    int ret = 0;
    while(f1 != f2)
    {
        if(dep[f1] < dep[f2])
        {
            swap(f1, f2);
            swap(u , v);
        }
        ret += query_sum(idx[f1], idx[u], 0, pos-1, 1);
        u = fa[f1]; f1 = top[u];
    }
    if(dep[u] < dep[v]) swap(u, v);
    ret += query_sum(idx[v], idx[u], 0, pos-1, 1);
    return ret;
}

int find_max(int u, int v)
{
    int f1 = top[u], f2 = top[v];
    int ret = -inf;
    while(f1 != f2)
    {
        if(dep[f1] < dep[f2])
        {
            swap(f1, f2);
            swap(u , v);
        }
        ret = max(ret, query_max(idx[f1], idx[u], 0, pos-1, 1) );
        u = fa[f1]; f1 = top[u];
    }
    if(dep[u] < dep[v]) swap(u, v);
    ret = max(ret, query_max(idx[v], idx[u], 0, pos-1, 1) );
    return ret;
}

int main()
{
    init();
    scanf("%d", &n);
    int u, v;
    for(int i=0; i<n-1; i++)
    {
        scanf("%d%d", &u, &v);
        add(u, v);
        add(v, u);
    }

    for(int i=1; i<=n; i++) scanf("%d", &w[i]);
    dfs1(1, -1, 0); dfs2(1, 1);

//    for(int i=1; i<=n; i++)
//    {
//        printf("son:%d size:%d dep:%d id:%d top:%d\n", son[i], size[i], dep[i], idx[i], top[i]);
//    }

    build(0, pos-1, 1);

    for(int i=1; i<=n; i++)
    {
        update(idx[i], w[i], 0, pos-1, 1);
    }

    int Q, a, b;
    char buf[10];
    cin >> Q;
    while(Q--)
    {
        scanf("%s%d%d", buf, &a, &b);
        if(buf[0] == 'C') update(idx[a], b, 0, pos-1, 1);
        else if(buf[1] == 'S') printf("%d\n", find_sum(a, b));
        else printf("%d\n", find_max(a, b));
    }
    return 0;
}