HYSBZ 1036 樹的統計Count（樹鏈剖分）

HYSBZ 1036 樹的統計Count

題目鏈接

就樹鏈剖分，線段樹維護sum和maxx即可

代碼：

#include 
#include 
#include 
#include 
using namespace std;

const int N = 30005;

int dep[N], fa[N], son[N], sz[N], top[N], id[N], idx, val[N];
int first[N], next[N * 2], vv[N * 2], en;

void init() {
	en = 0;
	idx = 0;
	memset(first, -1, sizeof(first));
}

void add_Edge(int u, int v) {
	vv[en] = v;
	next[en] = first[u];
	first[u] = en++;
}

void dfs1(int u, int f, int d) {
	dep[u] = d;
	sz[u] = 1;
	fa[u] = f;
	son[u] = 0;
	for (int i = first[u]; i + 1; i = next[i]) {
		int v = vv[i];
		if (v == f) continue;
		dfs1(v, u, d + 1);
		sz[u] += sz[v];
		if (sz[son[u]] < sz[v])
			son[u] = v;
	}
}

void dfs2(int u, int tp) {
	id[u] = ++idx;
	top[u] = tp;
	if (son[u]) dfs2(son[u], tp);
	for (int i = first[u]; i + 1; i = next[i]) {
		int v = vv[i];
		if (v == fa[u] || v == son[u]) continue;
		dfs2(v, v);
	}
}

const int INF = 0x3f3f3f3f;

int n;

#define lson(x) ((x<<1)+1)
#define rson(x) ((x<<1)+2)

struct Node {
	int l, r, maxx, sum;
} node[N * 4];

void pushup(int x) {
	node[x].maxx = max(node[lson(x)].maxx, node[rson(x)].maxx);
	node[x].sum = node[lson(x)].sum + node[rson(x)].sum;
}

void build(int l, int r, int x = 0) {
	node[x].l = l; node[x].r = r;
	if (l == r) {
		node[x].maxx = node[x].sum = val[l];
		return;
	}
	int mid = (l + r) / 2;
	build(l, mid, lson(x));
	build(mid + 1, r, rson(x));
	pushup(x);
}

void add(int v, int val, int x = 0) {
	if (node[x].l == node[x].r) {
		node[x].sum = node[x].maxx = val;
		return;
	}
	int mid = (node[x].l + node[x].r) / 2;
	if (v <= mid) add(v, val, lson(x));
	if (v > mid) add(v, val, rson(x));
	pushup(x);
}

int query(int l, int r, int ty, int x = 0) {
	if (node[x].l >= l && node[x].r <= r) {
		if (ty) return node[x].sum;
		return node[x].maxx;
	}
	int mid = (node[x].l + node[x].r) / 2;
	int ans = -INF, tmp;
	if (ty) ans = 0;
	if (l <= mid) {
		tmp = query(l, r, ty, lson(x));
		if (ty) ans += tmp;
		else ans = max(ans, tmp);
	}
	if (r > mid) {
		tmp = query(l, r, ty, rson(x));
		if (ty) ans += tmp;
		else ans = max(ans, tmp);
	}
	return ans;
}

int gao(int u, int v, int ty) {
	int ans = -INF, tmp; if (ty) ans = 0;
	int tp1 = top[u], tp2 = top[v];
	while (tp1 != tp2) {
		if (dep[tp1] < dep[tp2]) {
			swap(tp1, tp2);
			swap(u, v);
		}
		tmp = query(id[tp1], id[u], ty);
		if (ty) ans += tmp;
		else ans = max(ans, tmp);
		u = fa[tp1];
		tp1 = top[u];
	}
	if (dep[u] > dep[v]) swap(u, v);
	tmp = query(id[u], id[v], ty);
	if (ty) ans += tmp;
	else ans = max(ans, tmp);
	return ans;
}

int main() {
	while (~scanf("%d", &n)) {
		init();
		int u, v;
		for (int i = 1; i < n; i++) {
			scanf("%d%d", &u, &v);
			add_Edge(u, v);
			add_Edge(v, u);
		}
		dfs1(1, 0, 1);
		dfs2(1, 1);
		for (int i = 1; i <= n; i++) scanf("%d", &val[id[i]]);
		build(1, idx);
		int q;
		char Q[10]; int a, b;
		scanf("%d", &q);
		while (q--) {
			scanf("%s%d%d", Q, &a, &b);
			if (Q[0] == 'C') add(id[a], b);
			else {
				int ty = 0;
				if (Q[1] == 'S') ty = 1;
				printf("%d\n", gao(a, b, ty));
			}
		}
	}
	return 0;
}