Description
给定一棵有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
对于每个询问操作,输出一行答案。
Sample Input
6 5
2 2 1 2 1 1
1 2
1 3
2 4
2 5
2 6
Q 3 5
C 2 1 1
Q 3 5
C 5 1 2
Q 3 5
Sample Output
3
1
2
HINT
数N<=10^5,操作数M<=10^5,所有的颜色C为整数且在[0, 10^9]之间。
题解
首先处理树,使用树链剖分;线段树则维护区间颜色不同数量及左右端点的颜色。更新时若左边的右端点和右边的左端点颜色相同,则数量减1。
代码
/**************************************************************
Problem: 2243
User: xingsuixing
Language: C++
Result: Accepted
Time:5276 ms
Memory:31076 kb
****************************************************************/
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
using namespace std;
const int N = 200010;
struct Edge {
int v, next;
} edge[N];
int n, m, tot = 0, prelc, prerc, root, tail = 0;
int indexx = 0, fa[N], dep[N], top[N], siz[N], son[N], in[N], out[N], a[N], head[N], color[N];
void add(int u, int v) {
tot ++;
edge[tot].v = v;
edge[tot].next = head[u];
head[u] = tot;
}
struct Node {
int l, r, ls, rs;
int sum, flag, lc, rc;
} t[500010];
void dfs1(int u, int f, int d) {
siz[u] = 1, fa[u] = f, dep[u] = d;
for(int i = head[u]; i; i = edge[i].next) {
int v = edge[i].v;
if(v == f) continue ;
dfs1(v, u, d + 1);
siz[u] += siz[v];
if(son[u] == -1 || siz[v] > siz[son[u]])
son[u] = v;
}
}
void dfs2(int u, int tp) {
indexx ++;
in[u] = out[u] = indexx, top[u] = tp;
color[indexx] = a[u];
if(son[u] == -1) return ;
dfs2(son[u], tp);
for(int i = head[u]; i; i = edge[i].next) {
int v = edge[i].v;
if(v == fa[u] || v == son[u]) continue ;
dfs2(v, v);
}
out[u] = indexx;
}
void update(int num) {
int ls = num << 1, rs = num << 1 | 1;
t[num].lc = t[ls].lc;
t[num].rc = t[rs].rc;
t[num].sum = t[ls].sum + t[rs].sum;
if(t[ls].rc == t[rs].lc)
t[num].sum -= 1;
}
void build(int num, int l, int r) {
t[num].l = l, t[num].r = r;
if(l == r) {
t[num].sum = 1, t[num].flag = 0;
t[num].lc = t[num].rc = color[l];
return ;
}
int mid = (l + r) >> 1;
build(num << 1, l, mid);
build(num << 1 | 1, mid + 1, r);
update(num);
}
void pushdown(int num) {
int ls = num << 1, rs = num << 1 | 1;
if(t[num].flag) {
t[ls].sum = t[rs].sum = 1;
t[ls].lc = t[ls].rc = t[num].flag;
t[rs].lc = t[rs].rc = t[num].flag;
t[rs].flag = t[ls].flag = t[num].flag;
t[num].flag = 0;
}
}
void modify(int num, int L, int R, int c) {
int l = t[num].l, r = t[num].r;
if(L <= l && r <= R) {
t[num].flag = c;
t[num].sum = 1;
t[num].lc = c, t[num].rc = c;
return ;
}
pushdown(num);
int mid = (l + r) >> 1;
if(L <= mid)
modify(num << 1, L, R, c);
if(R > mid)
modify(num << 1 | 1, L, R, c);
update(num);
}
void delta(int u, int v, int c) {
int f1 = top[u], f2 = top[v];
while(f1 != f2){
if(dep[f1] < dep[f2]) swap(f1, f2), swap(u, v);
modify(root, in[f1], in[u], c);
u = fa[f1], f1 = top[u];
}
if(dep[u] > dep[v]) swap(u, v);
modify(root, in[u], in[v], c);
}
int query(int num, int pos, int val, int L, int R) {
int l = t[num].l, r = t[num].r;
if(l == L) prelc = t[num].lc;
if(r == R) prerc = t[num].rc;
if(L <= l && r <= R)
return t[num].sum;
pushdown(num);
int mid = (l + r) >> 1, rt = 0;
if(L > mid)
rt = query(num << 1 | 1, pos, val, L, R);
else if(R <= mid)
rt = query(num << 1, pos, val, L, R);
else {
rt = query(num << 1, pos, mid, L, R) + query(num << 1 | 1, mid + 1, val, L, R);
if(t[num << 1].rc == t[num << 1 | 1].lc)
rt -= 1;
}
return rt;
}
int check(int u, int v) {
int co1 = -1, co2 = -1;
int f1 = top[u], f2 = top[v], sum = 0;
while(f1 != f2) {
if(dep[f1] < dep[f2]) swap(f1, f2), swap(u, v), swap(co1, co2);
sum += query(1, in[f1], in[u], in[f1], in[u]);
if(prerc == co1) sum --;
co1 = prelc;
u = fa[f1];
f1 = top[u];
}
if(dep[u] < dep[v]) swap(u, v), swap(co1, co2);
sum += query(1, in[v], in[u], in[v], in[u]);
if(prerc == co1) sum --;
if(prelc == co2) sum --;
return sum;
}
int main(){
memset(son, -1, sizeof(son));
scanf("%d %d", &n, &m);
for(int i = 1; i <= n; i ++)
scanf("%d", &a[i]);
for(int i = 1; i < n; i ++) {
int u, v;
scanf("%d %d", &u, &v);
add(u, v), add(v, u);
}
dfs1(1, -1, 0); dfs2(1, 1);
int root = 1;
build(1, 1, n);
char s[2];
while(m --) {
scanf("%s", s);
if(s[0] == 'Q') {
int u, v;
scanf("%d %d", &u, &v);
printf("%d\n", check(u, v));
}
else {
int u, v, c;
scanf("%d %d %d", &u, &v, &c);
delta(u, v, c);
}
}
return 0;
}