1. 程式人生 > >2243 染色(樹剖+區間修改線段樹)

2243 染色(樹剖+區間修改線段樹)

給定一棵有n個節點的無根樹和m個操作,操作有2類: 1、將節點a到節點b路徑上所有點都染成顏色c; 2、詢問節點a到節點b路徑上的顏色段數量(連續相同顏色被認為是同一段), 如“112221”由3段組成:“11”、“222”和“1”。 請你寫一個程式依次完成這m個操作。 Input 第一行包含2個整數n和m,分別表示節點數和運算元; 第二行包含n個正整數表示n個節點的初始顏色 下面 行每行包含兩個整數x和y,表示x和y之間有一條無向邊。 下面 行每行描述一個操作: “C a b c”表示這是一個染色操作,把節點a到節點b路徑上所有點(包括a和b)都染成顏色c; “Q a b”表示這是一個詢問操作,詢問節點a到節點b(包括a和b)路徑上的顏色段數量。 Output 對於每個詢問操作,輸出一行答案。

幹,真的很難。 樹論的題程式碼一個比一個長,一個比一個難調。 wa了幾發才看見原來0也算一種顏色,趕緊把lazy的初始值改成-1。

#include<bits/stdc++.h>
using namespace std;

const int maxn = 1e5 + 5;

vector<int> G[maxn];
int n, m, in[maxn];
//子節點個數,深度
int sz[maxn], dep[maxn];
//重兒子,父節點
int ch[maxn], fa[maxn];
//重鏈開頭,dfs序
int top[maxn], tid[maxn], tid2[maxn];
int tot;


void dfs1(int u, int f, int d) {
	sz[u] = 1;
	fa[u] = f;
	dep[u] = d;
	for(int i = 0; i < G[u].size(); i++) {
		int v = G[u][i];
		if(v == f) {
			continue;
		}
		dfs1(v, u, d + 1);
		sz[u] += sz[v];
		if(ch[u] == -1 || sz[v] > sz[ch[u]]) {
			ch[u] = v;
		}
	}
}

void dfs2(int u, int tp) {
	top[u] = tp;
	tid[u] = ++tot;
	tid2[tot] = u;
	if(ch[u] == -1) {
		return;
	}
	dfs2(ch[u], tp);
	for(int i = 0; i < G[u].size(); i++) {
		int v = G[u][i];
		if(v != ch[u] && v != fa[u]) {
			dfs2(v, v);
		}
	}
}


struct SegmentTree {
	int sum[maxn << 2], lazy[maxn << 2];
	int left[maxn << 2], right[maxn << 2];

	void pushUp(int i) {
		left[i] = left[i << 1];
		right[i] = right[i << 1 | 1];
		sum[i] = sum[i << 1] + sum[i << 1 | 1] - (left[i << 1 | 1] == right[i << 1]);
	}

	void pushDown(int i) {
		if(lazy[i] == -1) {
			return;
		}
		left[i << 1] = right[i << 1] = lazy[i];
		left[i << 1 | 1] = right[i << 1 | 1] = lazy[i];
		lazy[i << 1] = lazy[i << 1 | 1] = lazy[i];
		sum[i << 1] = sum[i << 1 | 1] = 1;
		lazy[i] = -1;
	}

	void build(int i, int l, int r) {
		lazy[i] = -1;
		if(l == r) {
			left[i] = right[i] = in[tid2[l]];
			sum[i] = 1;
			return;
		}
		int mid = (l + r) >> 1;
		build(i << 1, l, mid);
		build(i << 1 | 1, mid + 1, r);
		pushUp(i);
	}

	void update(int i, int l, int r, int L, int R, int cor) {
		if(l >= L && r <= R) {
			sum[i] = 1;
			left[i] = right[i] = cor;
			lazy[i] = cor;
			return;
		}
		pushDown(i);
		int mid = (l + r) >> 1;
		if(L <= mid) {
			update(i << 1, l, mid, L, R, cor);
		}
		if(R > mid) {
			update(i << 1 | 1, mid + 1, r, L, R, cor);
		}
		pushUp(i);
	}

	int query(int i, int l, int r, int L, int R) {
		if(l >= L && r <= R) {
			return sum[i];
		}
		pushDown(i);
		int mid = (l + r) >> 1;
		if(R <= mid) {
			return query(i << 1, l, mid, L, R);
		}
		if(L > mid) {
			return query(i << 1 | 1, mid + 1, r, L, R);
		}
		return query(i << 1, l, mid, L, R) + query(i << 1 | 1, mid + 1, r, L, R) - (left[i << 1 | 1] == right[i << 1]);
	}

	int query(int i, int l, int r, int pos) {
		if(l == pos) {
			return left[i];
		}
		if(r == pos) {
			return right[i];
		}
		pushDown(i);
		int mid = (l + r) >> 1;
		if(pos <= mid) {
			return query(i << 1, l, mid, pos);
		}
		return query(i << 1 | 1, mid + 1, r, pos);
	}

	int solve(int u, int v, int val) {
		int f1 = top[u], f2 = top[v];
		// 不在同一條鏈上
		int ans = 0;
		while(f1 != f2) {
			if(dep[f1] < dep[f2]) {
				swap(f1, f2);
				swap(u, v);
			}
			if(val >= 0) {
				update(1, 1, n, tid[f1], tid[u], val);
			} else {
				ans += query(1, 1, n, tid[f1], tid[u]);
				ans -= query(1, 1, n, tid[fa[f1]]) == query(1, 1, n, tid[f1]);
			}
			u = fa[f1];
			f1 = top[u];
		}
		if(dep[u] < dep[v]) {
			swap(u, v);
		}
		if(val >= 0) {
			update(1, 1, n, tid[v], tid[u], val);
			return 0;
		}
		return ans + query(1, 1, n, tid[v], tid[u]);
	}

} st;

int main() {
	char op[3];
	scanf("%d%d", &n, &m);
	memset(ch, -1, sizeof(ch));
	tot = 0;
	for(int i = 1;i <= n; i++){
		G[i].clear();
	}
	
	for(int i = 1; i <= n; i++) {
		scanf("%d", &in[i]);
	}

	int u, v, val;
	for(int i = 1; i <= n - 1; i++) {
		scanf("%d%d", &u, &v);
		G[u].push_back(v);
		G[v].push_back(u);
	}

	dfs1(1, 0, 0);
	dfs2(1, 1);
	st.build(1, 1, n);

	while(m--) {
		scanf("%s%d%d", op, &u, &v);
		if(op[0] == 'C') {
			scanf("%d", &val);
			st.solve(u, v, val);
		} else {
			printf("%d\n", st.solve(u, v, -1));
		}
	}

	return 0;
}