【数据结构 树 树链剖分】luogu_2486 [SDOI2011]染色

题意

求树上路径点颜色的块数，带修改操作。

思路

先把树剖成若干链，用线段树维护区间的块数。

查询统计答案时，当一条链调到另一条链，判断一下这两个端点是否相等（即为同一块）。两个点跳时都这样操作。

当两个点到了同一条链上，需要两个端点都和它们上一次查询到的端点判断（因为在循环中去掉的是上一次和现在的重复，此处是循环结束后特判）。

细节详见代码。

代码

#include <cstdio>
#include <algorithm>

struct segmentTree {
	int ls, rs, l, r, lc, rc, dat, lazy;
}t[400001];
int father[100001], dep[100001], size[100001], son[100001], rev[100001], top[100001], seg[100001];
int head[100001], ver[200001], next[200001];
int a[100001];
int n, m, tot, cnt;

void add(int u, int v) {
	ver[++tot] = v;
	next[tot] = head[u];
	head[u] = tot;
}

void dfs1(int u, int fa) {
	father[u] = fa;
	size[u] = 1;
	dep[u] = dep[fa] + 1;
	for (int i = head[u]; i; i = next[i]) {
		if (ver[i] == fa) continue;
		dfs1(ver[i], u);
		size[u] += size[ver[i]];
		if (size[ver[i]] > size[son[u]]) son[u] = ver[i];
	}
}

void dfs2(int u, int t) {
	top[u] = t;
	seg[u] = ++cnt;
	rev[cnt] = u;
	if (!son[u]) return;
	dfs2(son[u], t);
	for (int i = head[u]; i; i = next[i]) {
		if (top[ver[i]]) continue;
		dfs2(ver[i], ver[i]);
	}
}

void build(int p, int l, int r) {
	t[p].l = l, t[p].r = r;
	if (l == r) {
		t[p].lc = t[p].rc = a[rev[l]];
		t[p].dat = 1;
		return;
	}
	int mid = l + r >> 1;
	build(t[p].ls = ++cnt, l, mid);
	build(t[p].rs = ++cnt, mid + 1, r);
	t[p].dat = t[t[p].ls].dat + t[t[p].rs].dat - (t[t[p].ls].rc == t[t[p].rs].lc);
	t[p].lc = t[t[p].ls].lc;
	t[p].rc = t[t[p].rs].rc;
}

void spread(int p) {
	if (!t[p].lazy) return;
	t[t[p].ls].lazy = t[t[p].rs].lazy = t[t[p].ls].lc = t[t[p].ls].rc = t[t[p].rs].lc = t[t[p].rs].rc = t[p].lazy;
	t[t[p].ls].dat = t[t[p].rs].dat = 1;
	t[p].lazy = 0;
}

int query(int p, int l, int r) {
	if (l <= t[p].l && t[p].r <= r)
		return t[p].dat;
	spread(p);
	int mid = t[p].l + t[p].r >> 1;
	if (l <= mid && r > mid) return query(t[p].ls, l, r) + query(t[p].rs, l, r)  - (t[t[p].ls].rc == t[t[p].rs].lc);
	else if (l <= mid) return query(t[p].ls, l, r);
	else return query(t[p].rs, l, r);
}

int queryc(int p, int pos) {
	if (t[p].l == t[p].r)
		return t[p].lc;
	spread(p);
	int mid = t[p].l + t[p].r >> 1;
	if (pos <= mid) return queryc(t[p].ls, pos);
	else return queryc(t[p].rs, pos);
}

int ask(int x, int y) {
	int res = 0, ans1 = -1, ans2 = -1, L = 0, R = 0;
	while (top[x] != top[y]) {
		if (dep[top[x]] < dep[top[y]]) std::swap(x, y), std::swap(ans1, ans2);
		res += query(1, seg[top[x]], seg[x]);
		L = queryc(1, seg[top[x]]);
		R = queryc(1, seg[x]);
		if (R == ans1) res--;
		ans1 = L;
		x = father[top[x]];
	}
	if (dep[x] > dep[y]) std::swap(x, y), std::swap(ans1, ans2);
	res += query(1, seg[x], seg[y]);
	L = queryc(1, seg[x]);
	R = queryc(1, seg[y]);
	if (ans2 == R) res--;
	if (ans1 == L) res--;
	return res;
}

void update(int p, int l, int r, int val) {
	if (l <= t[p].l && t[p].r <= r) {
		t[p].dat = 1;
		t[p].lazy = t[p].lc = t[p].rc = val;
		return;
	}
	spread(p);
	int mid = t[p].l + t[p].r >> 1;
	if (l <= mid) update(t[p].ls, l, r, val);
	if (r > mid) update(t[p].rs, l, r, val);
	t[p].dat = t[t[p].ls].dat + t[t[p].rs].dat - (t[t[p].ls].rc == t[t[p].rs].lc);
	t[p].lc = t[t[p].ls].lc;
	t[p].rc = t[t[p].rs].rc;
}

void modify(int x, int y, int val) {
	while (top[x] != top[y]) {
		if (dep[top[x]] < dep[top[y]]) std::swap(x, y);
		update(1, seg[top[x]], seg[x], val);
		x = father[top[x]];
	}
	if (dep[x] > dep[y]) std::swap(x, y);
	update(1, seg[x], seg[y], val);
}

int main() {
	scanf("%d %d", &n, &m);
	for (int i = 1; i <= n; i++)
		scanf("%d", &a[i]);
	for (int i = 1, x, y; i < n; i++)
		scanf("%d %d", &x, &y), add(x, y), add(y, x);
	dfs1(1, 0);
	dfs2(1, 1);
	build(cnt = 1, 1, n);
	for (int a, b, c; m; m--) {
		char op;
		for (op = getchar(); op != 'Q' && op != 'C'; op = getchar());
		if (op == 'Q') {
			scanf("%d %d", &a, &b);
			printf("%d\n", ask(a, b));
		} else {
			scanf("%d %d %d", &a, &b, &c);
			modify(a, b, c);
		}
	}
}