HYSBZ 2243 染色

题目链接：http://www.lydsy.com/JudgeOnline/problem.php?id=2243

题意:

给定一棵有n个节点的无根树和m个操作，操作有2类：

1、将节点a到节点b路径上所有点都染成颜色c；

2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），如“112221”由3段组成：“11”、“222”和“1”。

思路：

权值在点上的树链剖分+线段树成段更新+线段树成段询问。

维护颜色端的数量方法：

如果左儿子区间的最右端颜色与右儿子区间的最左端区间相同。

sum[rt] = sum[rt<<1]+sum[rt<<1|1] - 1;

否则sum[rt] = sum[rt<<1]+sum[rt<<1|1];

树链剖分部分也进行相应维护。

#include <iostream>
#include <cstring>
#include <string>
#include <cstdio>
#include <vector>
using namespace std;
int n, q;
#define maxn 100010
#define lson l, m, rt<<1
#define rson m+1, r, rt<<1|1
int siz[maxn], top[maxn], fa[maxn], son[maxn], dep[maxn];
int w[maxn], fw[maxn];
int A[maxn];
vector <int> mp[maxn];
int pos;
int Left, Right;
int sum[maxn<<2], cl[maxn<<2], cr[maxn<<2], col[maxn<<2];
void PushUp(int rt)
{
if(cr[rt<<1] == cl[rt<<1|1]) sum[rt] = sum[rt<<1] + sum[rt<<1|1] - 1;
else sum[rt] = sum[rt<<1] + sum[rt<<1|1];
cl[rt] = cl[rt<<1];
cr[rt] = cr[rt<<1|1];
}
void PushDown(int rt, int m)
{
if(col[rt] >= 0)
{
col[rt<<1] = col[rt<<1|1] = col[rt];
sum[rt<<1] = sum[rt<<1|1] = 1;
cl[rt<<1] = cl[rt<<1|1] = cr[rt<<1] = cr[rt<<1|1] = cl[rt];
col[rt] = -1;
}
}
void build(int l, int r, int rt)
{
col[rt] = -1;
if(l == r)
{
cl[rt] = cr[rt] = A[fw[l]];
sum[rt] = 1;
return;
}
int m = (l+r)>>1;
build(lson);
build(rson);
PushUp(rt);
}
void update(int L, int R, int c, int l, int r, int rt)
{
if(L <= l && R >= r)
{
col[rt] = c; sum[rt] = 1; cl[rt] = c; cr[rt] = c;
return;
}
PushDown(rt, r-l+1);
int m = (l+r)>>1;
if(L <= m) update(L, R, c, lson);
if(R > m) update(L, R, c, rson);
PushUp(rt);
}
int query(int L, int R, int l, int r, int rt)
{
if(l == L) Left = cl[rt];
if(r == R) Right = cr[rt];
if(L <= l && R >= r)
{
return sum[rt];
}
PushDown(rt, r-l+1);
int m = (l+r)>>1;
int ret = 0;
//合并线段树 这里也要修改!
if(R <= m) return query(L, R, lson);
else if(L > m) return query(L, R, rson);
else
{
ret = query(L,R,lson) + query(L, R, rson);
if(cl[rt<<1|1] == cr[rt<<1]) ret --;
return ret;
}
//if(L <= m) ret += query(L, R, lson);
//if(R > m) ret += query(L, R, rson);
}
int dfs1(int u, int pre, int deep)
{
siz[u] = 1; dep[u] = deep; fa[u] = pre;
int mmax = 0;
for(int i = 0; i < mp[u].size(); i++)
{
if(mp[u][i] != pre)
{
int temp = dfs1(mp[u][i], u, deep+1);
siz[u] += temp;
if(son[u] == -1 || temp >= mmax)
{
son[u] = mp[u][i];
mmax = temp;
}
}
}
return siz[u];
}
void dfs2(int u, int val)
{
top[u] = val;
if(son[u] != -1)
{
w[u] = ++pos;
fw[w[u]] = u;
dfs2(son[u], val);
}
else if(son[u] == -1)
{
w[u] = ++pos;
fw[w[u]] = u;
return;
}
for(int i = 0; i < mp[u].size(); i++)
{
if(mp[u][i] != son[u] && mp[u][i] != fa[u]) dfs2(mp[u][i], mp[u][i]);
}
}
int find(int u, int v)
{
int mark1 = -1, mark2 = -1;
int f1 = top[u], f2 = top[v];
int temp = 0; int precolor = -1;
while(f1 != f2)
{
if(dep[f1] < dep[f2])   //规定f1更深
{
swap(f1, f2);
swap(u, v);
swap(mark1, mark2);  //交替
}
temp += query(w[f1], w[u], 1, pos, 1);
if(Right == mark1) temp --;
mark1 = Left;
u = fa[f1]; f1 = top[u];
}
if(dep[u] < dep[v])  //规定u更深
{
swap(u, v);
swap(mark1, mark2);
}
temp += query(w[v], w[u], 1, pos, 1);
if(Right == mark1) temp--;
if(Left == mark2) temp--;
return temp;
}
void UPDATE(int u, int v, int color)
{
int f1 = top[u], f2 = top[v];
while(f1 != f2)
{
if(dep[f1] < dep[f2])
{
swap(f1, f2);
swap(u, v);
}
update(w[f1], w[u], color, 1, pos, 1);
u = fa[f1]; f1 = top[u];
}
if(dep[u] > dep[v]) swap(u, v);
update(w[u], w[v], color, 1, pos, 1);
}
int main()
{
//freopen("in.txt", "r", stdin);
while(~scanf("%d%d", &n, &q))
{
for(int i = 1; i <= n; i++) mp[i].clear();
pos = 0;
memset(son, -1, sizeof(son));
for(int i = 1; i <= n; i++) scanf("%d", &A[i]);
for(int i = 1; i <= n-1; i++)
{
int a, b; scanf("%d%d", &a, &b);
mp[a].push_back(b);
mp[b].push_back(a);
}
dfs1(1, -1, 1);
dfs2(1, 1);
build(1, pos, 1);

char op[5];
while(q--)
{
scanf("%s", op);
int a, b, c;
if(op[0] == 'Q')
{
scanf("%d%d", &a, &b);
printf("%d\n", find(a, b));
}
else if(op[0] == 'C')
{
scanf("%d%d%d", &a, &b, &c);
UPDATE(a, b, c);
}
}
}
return 0;
}