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;
}