HYSBZ2243 染色【树链剖分】

题意：给定一棵树，有两种操作，1.问两点之间点的颜色段数量 2.把两点之间点都染成某颜色

思路：颜色段，线段树维护区间左端点和右端点，区间合并的时候，两区间中间的两个点颜色相同，就减1。

树链剖分查询，重链跳转那里，看看当前点 和 跳上去的点颜色是否相同，相同减1。

#include<stdio.h>
#include<iostream>
#include<string.h>
#include<string>
#include<stdlib.h>
#include<math.h>
#include<vector>
#include<list>
#include<map>
#include<stack>
#include<queue>
#include<algorithm>
#include<numeric>
#include<functional>
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
#define lson(x) 2*x
#define rson(x) 2*x+1
const int maxn = 2e5+5;
struct Edge
{
	int to,next;
}edge[maxn*2];
int a[maxn]; 
int head[maxn],tot,cnt;
int root,size[maxn],dep[maxn],fa[maxn],son[maxn];
int top[maxn],id[maxn],val[maxn];
/*
size -> 以x为根节点子树的大小   dep -> 深度 
fa   -> 父节点                  son -> 重儿子 
top  -> x所在重链的顶端节点     
id -> 这个点在线段树中的位置 
*/
struct data
{
	int l,r,num,laz,zuo,you;
}node[maxn*4];

void init()
{
	memset(head,-1,sizeof head);
	tot = 0;
	cnt = 0;
}

void add(int x,int y)
{
	edge[tot].to = y;
	edge[tot].next = head[x];
	head[x] = tot++;
}

int dfs1(int x,int pre,int d)
{
	size[x] = 1;
	fa[x] = pre;
	dep[x] = d;
	son[x] = 0;
	for(int i = head[x]; i != -1; i = edge[i].next)
	{
		int y = edge[i].to;
		if(y != pre)
		{
			size[x] += dfs1(y,x,d+1);
			if(size[y] > size[son[x]])
				son[x] = y;
		}
	}
	return size[x];
}

void dfs2(int x,int pre)
{
	id[x] = ++cnt;
	top[x] = pre;
	if(son[x]) dfs2(son[x],pre); //重链 
	for(int i = head[x]; i != -1; i = edge[i].next)
	{
		int y = edge[i].to;
		if(y != fa[x] && y != son[x])
			dfs2(y,y); //轻链
	}
}

void pushup(int cnt)
{
	node[cnt].num = node[lson(cnt)].num + node[rson(cnt)].num;
	if(node[lson(cnt)].you == node[rson(cnt)].zuo)
		node[cnt].num--;
	node[cnt].zuo = node[lson(cnt)].zuo;
	node[cnt].you = node[rson(cnt)].you;
}

void build(int x,int y, int cnt)
{
	node[cnt].l = x;
	node[cnt].r = y;
	node[cnt].laz = -1;
	if(x == y)
	{
		node[cnt].num = 1;
		node[cnt].zuo = node[cnt].you = val[x];
		return;
	}
	int mid = (x+y) / 2;
	build(x,mid,lson(cnt));
	build(mid+1,y,rson(cnt));
	pushup(cnt);
}

void down(int cnt)
{
	if(node[cnt].laz == -1)
		return;
	int val = node[cnt].laz;
	node[lson(cnt)].laz = val;
	node[rson(cnt)].laz = val;
	node[lson(cnt)].zuo = node[lson(cnt)].you = val;
	node[rson(cnt)].zuo = node[rson(cnt)].you = val;
	node[lson(cnt)].num = 1;
	node[rson(cnt)].num = 1;
	node[cnt].laz = -1;
}

void up(int x,int y,int cnt,int val)
{
	if(x == node[cnt].l && y == node[cnt].r)
	{
		node[cnt].laz = val;
		node[cnt].zuo = node[cnt].you = val;
		node[cnt].num = 1;
		return;
	}
	down(cnt);
	int fa = 2*cnt;
	if(x <= node[fa].r)
	{
		if(y <= node[fa].r)
			up(x,y,fa,val);
		else
			up(x,node[fa].r,fa,val);
	}
	fa++;
	if(y >= node[fa].l)
	{
		if(x >= node[fa].l)
			up(x,y,fa,val);
		else
			up(node[fa].l,y,fa,val);
	}
	pushup(cnt);
}

int yan(int x,int y,int cnt)
{
	if(x == node[cnt].l && y == node[cnt].r)
	{
		return node[cnt].zuo;
	}
	if(node[cnt].laz != -1)
		return node[cnt].laz;
	int fa = 2*cnt;
	if(x <= node[fa].r)
	{
		if(y <= node[fa].r)
			return yan(x,y,fa);
		else
			return yan(x,node[fa].r,fa);
	}
	fa++;
	if(y >= node[fa].l)
	{
		if(x >= node[fa].l)
			return yan(x,y,fa);
		else
			return yan(node[fa].l,y,fa);
	}
}

data meg(data a, data b)
{
	data res;
	res.num = a.num + b.num;
	if(a.you == b.zuo)
		res.num--;
	res.zuo = a.zuo;
	res.you = b.you;
	return res;
}

data fid(int x,int y,int cnt)
{
	down(cnt);
	int res = 0;
	if(x == node[cnt].l && y == node[cnt].r)
		return node[cnt];
	int mid = (node[cnt].l + node[cnt].r) / 2;
	if(y <= mid)
		return fid(x,y,lson(cnt));
	else if(x >= mid+1)
		return fid(x,y,rson(cnt));
	else
		return meg( fid(x,mid,lson(cnt)) , fid(mid+1,y,rson(cnt)) );	
}

int qurey(int u,int v)
{
	int ans = 0;
	data res;
	while(top[u] != top[v])
	{
		if(dep[top[u]] < dep[top[v]])
			swap(u,v);
		res = fid(id[top[u]],id[u],1);
		ans += res.num;
		u = top[u];
		if(yan(id[u],id[u],1) == yan(id[fa[u]],id[fa[u]],1))
			ans--;
		u = fa[u];
	}
	if(dep[u] > dep[v])
		swap(u,v);
	res = fid(id[u],id[v],1);
	ans += res.num;
	return ans;//边权,存的都是点的父边，最上面的点用son[u] 
}

void gai(int u,int v,int val)
{
	while(top[u] != top[v])
	{
		if(dep[top[u]] < dep[top[v]])
			swap(u,v);
		up(id[top[u]],id[u],1,val);
		u = top[u];
		u = fa[u];
	}
	if(dep[u] > dep[v])
		swap(u,v);
	up(id[u],id[v],1,val);
}

int main()
{
	int n,q;
	while(scanf("%d%d",&n,&q)!=EOF)
	{
		init();
		for(int i = 1; i <= n; i++)
			scanf("%d",&a[i]);
		for(int i = 1; i < n; i++)
		{
			int a,b;
			scanf("%d%d",&a,&b);
			add(a,b);
			add(b,a);
		}
		root = (n+1) / 2;
		dfs1(root,root,1);
		dfs2(root,root);
		for(int i = 1; i <= n; i++) //点权 
			val[id[i]] = a[i];		
		build(1,cnt,1);
		char op[10];
		while(q--)
		{
			scanf("%s",op);
			if(op[0] == 'Q')
			{
				int a,b;
				scanf("%d%d",&a,&b);
				printf("%d\n",qurey(a,b));
			}
			else
			{
				int a,b,c;
				scanf("%d%d%d",&a,&b,&c);
				gai(a,b,c);
			}
		}	
	}
}