bzoj 2243 luogu p2486 树链剖分+线段树

传送门

题意：中文题。

思路：进行树链剖分后用线段树维护区间，注意区间合并时相邻端点颜色相同时ans-1，我刚开始对区间端点的处理是建立在原数组上的，改成结构体储存区间左右端点颜色就A了。

代码：

#include <iostream>

using namespace std;

typedef long long ll;
const int N = 3e5 + 10;

struct edge
{
	int to, next;
}Edge[N<<1];

struct node
{
	int sum, lc, rc;
}Node[N<<2];

int n, m, a[N], head[N], lazy[N<<2];
int cnt, f[N], dep[N], size[N], son[N], rk[N], top[N], id[N];
//树链剖分模板
void add_edge(int A,int B)
{
	Edge[++cnt].to = B;
	Edge[cnt].next = head[A];
	head[A] = cnt;
}

void dfs1(int v,int fa,int depth)
{
	f[v] = fa;
	dep[v] = depth;
	size[v] = 1;
	for(int i = head[v]; i; i = Edge[i].next){
		int to = Edge[i].to;
		if(to == fa) continue;
		dfs1(to,v,depth+1);
		size[v] += size[to];
		if(size[to] > size[son[v]])
			son[v] = to;
	}
}

void dfs2(int v,int tp)
{
	top[v] = tp;
	id[v] = ++cnt;
	rk[cnt] = v;
	if(!son[v]) return ;
	dfs2(son[v],tp);
	for(int i = head[v]; i; i = Edge[i].next){
		int to = Edge[i].to;
		if(to != son[v] && to != f[v])
			dfs2(to,to);
	}
}

void pushUp(int rt)
{
	Node[rt].sum = Node[rt<<1].sum + Node[rt<<1|1].sum - (Node[rt<<1].rc == Node[rt<<1|1].lc);
	Node[rt].lc = Node[rt<<1].lc, Node[rt].rc = Node[rt<<1|1].rc;
}

void pushDown(int rt)
{
	if(lazy[rt] == -1) return ;
	Node[rt<<1].sum = Node[rt<<1|1].sum = 1;
	Node[rt<<1].lc = Node[rt<<1].rc = Node[rt<<1|1].lc = Node[rt<<1|1].rc = lazy[rt];
	lazy[rt<<1] = lazy[rt<<1|1] = lazy[rt];
	lazy[rt] = -1;
}

void build(int l,int r,int rt)
{
	lazy[rt] = -1;
	if(l == r){
		Node[rt].sum = 1;
		Node[rt].lc = Node[rt].rc = a[rk[l]];
		return ;
	}
	int m = (l+r)>>1;
	build(l,m,rt<<1);
	build(m+1,r,rt<<1|1);
	pushUp(rt);
}

void update(int L,int R,int c,int l,int r,int rt)
{
	if(L <= l && r <= R){
		Node[rt].sum = 1;
		Node[rt].lc = Node[rt].rc = c;
		lazy[rt] = c;
		return ;
	}
	int m = (l+r)>>1;
	pushDown(rt);
	if(L <= m) update(L,R,c,l,m,rt<<1);
	if(m <  R) update(L,R,c,m+1,r,rt<<1|1);
	pushUp(rt);
}

int query(int L,int R,int l,int r,int rt)
{
	if(L <= l && r <= R){
		return Node[rt].sum;
	}
	int m = (l+r)>>1;
	ll ans = 0, flag = 0;
	pushDown(rt);
	if(L <= m) ans += query(L,R,l,m,rt<<1), flag += 1;
	if(m <  R) ans += query(L,R,m+1,r,rt<<1|1), flag += 1;
	if(flag == 2) ans -= (Node[rt<<1].rc == Node[rt<<1|1].lc);
	return ans;
}

int getcolor(int L,int l,int r,int rt)
{
	if(l == r) return Node[rt].lc;
	pushDown(rt);
	int m = (l+r)>>1;
	if(L <= m) return getcolor(L,l,m,rt<<1);
	else	   return getcolor(L,m+1,r,rt<<1|1);
}

void updates(int A,int B,int C)
{
	while(top[A] != top[B]){
		if(dep[top[A]] < dep[top[B]]) swap(A,B);
		update(id[top[A]],id[A],C,1,n,1);
		A = f[top[A]];
	}
	if(id[A] > id[B]) swap(A,B);
	update(id[A],id[B],C,1,n,1);
}

ll querys(int A,int B)
{
	ll ans = 0;
	while(top[A] != top[B]){
		if(dep[top[A]] < dep[top[B]]) swap(A,B);
		ans += query(id[top[A]],id[A],1,n,1);
		if(getcolor(id[top[A]],1,n,1) == getcolor(id[f[top[A]]],1,n,1)) ans -= 1;
		A = f[top[A]]; //在这里对所有lazy标记进行下移，如果该链的头节点到它的父节点的颜色相同时，答案减一
	}
	if(id[A] > id[B]) swap(A,B);
	return ans + query(id[A],id[B],1,n,1);
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0), cout.tie(0);
	int A, B, C;
	char oper;
	cin>>n>>m;
	for(int i = 1; i <= n; ++i)
		cin>>a[i];
	for(int i = 0; i < n - 1; ++i){
		cin>>A>>B;
		add_edge(A,B); add_edge(B,A);
	}
	cnt = 0;
	dfs1(1,0,1);
	dfs2(1,1);
	build(1,n,1);
	while(m--){
		cin>>oper;
		if(oper == 'Q'){
			cin>>A>>B;
			cout<<querys(A,B)<<'\n';
		}
		else{
			cin>>A>>B>>C;
			updates(A,B,C);
		}
	}
	return 0;
}