1036: [ZJOI2008]树的统计Count 树链剖分裸题

题目链接：点击打开链接

树链剖分裸题：

#include<stdio.h>
#include<string.h>
#include<iostream>
#include<stdlib.h>
#define N 60003
#define L(x) (x<<1)
#define R(x) (x<<1|1)
#define Mid(x,y) ((x+y)>>1)
#define inf 10000000
using namespace std;

struct Edge{
	int to, nex;
}edge[N*2];
int head[N], edgenum;
void add(int u, int v){
	Edge E={v,head[u]}; 	edge[edgenum] = E; 	head[u] = edgenum++;
}

int n, Query, a[N];

int fa[N];//父节点
int dep[N];//深度
int son[N];//重儿子
int size[N];//子树节点数
int p[N]; //p[v]表示v在线段树中的位置
int fp[N]; //与p数组相反
int top[N]; // top[v] 表示v所在的重链的顶端节点 (这样 v与top[v] 的重路径可以直接在线段树上操作)lkjiyhtrf saQ

int dfs(int u, int Father, int deep){
	fa[u] = Father;		size[u] = 1;	dep[u] = deep;
	for(int i = head[u]; ~i; i = edge[i].nex){
		int v = edge[i].to;
		if(v == Father)continue;
		size[u] += dfs(v, u, deep+1);
		if(son[u] == -1 || size[v] > size[son[u]])son[u] = v;
	}
	return size[u];
}

int tree_id; //tree_id表示线段树中所有的边（给每条边编号）
void Have_p(int u, int Father){
	top[u] = Father;
	p[u] = tree_id++; fp[p[u]] = u;
	if(son[u] == -1)return ;
	Have_p(son[u], top[u]); //注意这里
	for(int i = head[u]; ~i; i = edge[i].nex){
		int v = edge[i].to;
		if(v != son[u] && v != fa[u])Have_p(v, v); //注意这里
	}
}

struct node{
	int l, r;
	int Max, Sum;
}tree[N*8];
void push_up(int id){
	tree[id].Max = max(tree[L(id)].Max, tree[R(id)].Max);
	tree[id].Sum = tree[L(id)].Sum + tree[R(id)].Sum;
}
void build(int l, int r, int id){
	tree[id].l = l, tree[id].r = r;
	if(l == r)
	{
		tree[id].Sum = tree[id].Max = a[fp[l]]; //注意这里
		return ;
	}
	int mid = Mid(l, r);
	build(l, mid, L(id));
	build(mid+1, r, R(id));
	push_up(id);
}

void updata(int pos, int val, int id){
	if(tree[id].l == tree[id].r)
	{
		tree[id].Max = tree[id].Sum = val;
		return ;
	}
	int mid = Mid(tree[id].l, tree[id].r);
	if(pos <= mid)updata(pos, val, L(id));
	else updata(pos, val, R(id));
	push_up(id);
}
int querySum(int l, int r, int id){
	if(l == tree[id].l && tree[id].r == r)return tree[id].Sum;
	int mid = Mid(tree[id].l, tree[id].r);
	if(r<=mid) return querySum(l, r, L(id));
	else if(mid<l)return querySum(l, r, R(id));
	return querySum(l, mid, L(id)) + querySum(mid+1, r, R(id));
}
int queryMax(int l, int r, int id){
	if(l == tree[id].l && tree[id].r == r)return tree[id].Max;
	int mid = Mid(tree[id].l, tree[id].r);
	if(r<=mid) return queryMax(l, r, L(id));
	else if(mid<l)return queryMax(l, r, R(id));
	return max(queryMax(l, mid, L(id)), queryMax(mid+1, r, R(id)));
}
int findSum(int u, int v){
	int f1 = top[u], f2 = top[v];
	int tmp = 0;
	while(f1 != f2)//相等就说明两点在同一重路径上 || 相同点
	{
		if(dep[f1]<dep[f2]) swap(f1, f2), swap(u, v); //注意这里
		tmp += querySum(p[f1], p[u], 1); //每次只爬其中一条重链！
		u = fa[f1]; //注意这里
		f1 = top[u];
	}
	if(dep[u] > dep[v])swap(u, v); //注意这里
	return tmp + querySum(p[u], p[v], 1);
}
int findMax(int u, int v){
	int f1 = top[u], f2 = top[v];
	int tmp = -inf;
	while(f1 != f2)
	{
		if(dep[f1]<dep[f2]) swap(f1, f2), swap(u, v);
		tmp = max(tmp, queryMax(p[f1], p[u], 1));
		u = fa[f1];
		f1 = top[u];
	}
	if(dep[u] > dep[v])swap(u, v);
	return max(tmp, queryMax(p[u], p[v], 1));
}

void init(){
	memset(head, -1, sizeof(head));
	memset(son, -1, sizeof(son));  //注意这里son初始化
	edgenum = 0;
	tree_id = 1; //这个的所有可用编号就是线段树的区间 注意这里是从1开始的
}
char op[10];
int main(){
	int u, v, i;
	while(~scanf("%d",&n)){
		init();
		for(i = 0; i < n-1; i++){scanf("%d %d",&u,&v);	add(u,v), add(v,u);	}
		for(i = 1; i <= n; i++)scanf("%d",&a[i]);
		dfs(1, 1, 0);
		Have_p(1, 1);
		build(1, n, 1);
		scanf("%d",&Query);
		while(Query--){
			scanf("%s%d%d",op,&u,&v);
			if(op[0] == 'C')
				updata(p[u],v,1); //注意这里
			else if(op[1] == 'M')
				printf("%d\n", findMax(u, v));
			else 
				printf("%d\n", findSum(u, v));
		}
	}
	return 0;
}

/*
4
1 2
2 3
4 1

4 2 1 3

12
QMAX 3 4
QMAX 3 3
QMAX 3 2
QMAX 2 3
QSUM 3 4
QSUM 2 1
CHANGE 1 5
QMAX 3 4
CHANGE 3 6
QMAX 3 4
QMAX 2 4
QSUM 3 4

1
5
QMAX 1 1
QSUM 1 1
CHANGE 1 4 
QMAX 1 1
QSUM 1 1


*/