题意:
给定一棵有n个节点的无根树和m个操作,操作有2类:
1)将节点a到节点b路径上所有点都染成颜色c;
2)询问节点a到节点n路径上的颜色段数量(连续相同颜色被认为是同一段),如“112221”由3段组成:“11”、“222”和“1”。
n, m <= 100000
思路:
树的形态没有改变,用树链剖分维护即可。
在线段树上维护区间内的颜色前缀,后缀以及颜色段数构成一个三元组(pre,suf,tot),则可很方便地合并信息:
(pre1, suf1, tot1) + (pre2, suf2, tot2) = (pre1, suf2, tot1 + tot2 - (suf1 == pre2))
注意一开始输入的颜色对应的点在线段树中的编号与原编号不同,查询操作合并信息要特别注意不要加反了。
又复习了一遍树剖,感觉细节比较多,写的时候要想清楚。。。
#include <cstdio>
#include <vector>
#define For(i,j,k) for(int i = j;i <= (int)k;i++)
#define lc (h << 1), L, M
#define rc (h << 1 | 1), M+1, R
#define root 1, 1, n
#define M ((L + R) >> 1)
using namespace std;
int n, m;
const int N = 100010;
int fa[N], id[N], top[N], co[N], c[N], sz[N], son[N], dep[N], e;
vector<int> G[N];
void DFS1(int h){
sz[h] = 1, dep[h] = dep[fa[h]] + 1;
For(i,0,G[h].size() - 1){
int v = G[h][i];
if(v == fa[h]) continue;
fa[v] = h;
DFS1(v);
sz[h] += sz[v];
if(sz[v] > sz[son[h]]) son[h] = v;
}
}
void DFS2(int h){
id[h] = ++e;
c[e] = co[h];
top[h] = son[fa[h]] == h ? top[fa[h]] : h;
if(son[h]) DFS2(son[h]);
For(i,0,G[h].size() - 1){
int v = G[h][i];
if(v == fa[h] || v == son[h]) continue;
DFS2(v);
}
}
struct info{
int pre, suf, tot;
info rev(){
return (info){suf, pre, tot};
}
info operator + (const info& A) const{
return (info){pre, A.suf, tot + A.tot - (suf == A.pre)};
}
};
struct SegmentTree{
info A[N<<2];
int col[N<<2];
void pushup(int h){
A[h] = A[h << 1] + A[h << 1 | 1];
}
void pushdown(int h, int L, int R){
if(col[h]){
Modify(lc, L, M, col[h]), Modify(rc, M+1, R, col[h]);
col[h] = 0;
}
}
void Create(int h, int L, int R){
if(L == R) A[h] = (info){c[L], c[L], 1}, col[h] = c[L];
else{
Create(lc), Create(rc);
pushup(h);
}
}
void Modify(int h, int L, int R, int l, int r, int v){
if(l <= L && R <= r)
A[h] = (info){v, v, 1}, col[h] = v;
else{
pushdown(h, L, R);
if(l <= M) Modify(lc, l, r, v);
if(r > M) Modify(rc, l, r, v);
pushup(h);
}
}
info Query(int h, int L, int R, int l, int r){
if(l <= L && R <= r) return A[h];
else{
pushdown(h, L, R);
if(r <= M) return Query(lc, l, r);
else if(l > M) return Query(rc, l, r);
else return Query(lc, l, r) + Query(rc, l, r);
}
}
}T;
void change(int u, int v, int val){
while(top[u] != top[v]){
if(dep[top[u]] < dep[top[v]]) swap(u, v);
T.Modify(1, 1, n, id[top[u]], id[u], val);
u = fa[top[u]];
}
if(dep[u] > dep[v]) swap(u, v);
T.Modify(root, id[u], id[v], val);
}
int query(int u, int v){
info l, r;
l = r = (info){0, 0, 0};
int x = u, y = v;
while(top[x] != top[y]){
if(dep[top[x]] < dep[top[y]]) y = fa[top[y]];
else x = fa[top[x]];
}
while(u != x) l = T.Query(root, id[top[u]], id[u]) + l, u = fa[top[u]];
while(v != y) r = T.Query(root, id[top[v]], id[v]) + r, v = fa[top[v]];
if(dep[u] > dep[v]) return (l.rev() + T.Query(1, 1, n, id[v], id[u]).rev() + r).tot;
else return (l.rev() + T.Query(root, id[u], id[v]) + r).tot;
}
int main(){
scanf("%d%d", &n, &m);
For(i,1,n) scanf("%d", &co[i]);
For(i,2,n){
int x, y;
scanf("%d%d", &x, &y);
G[x].push_back(y), G[y].push_back(x);
}
DFS1(1), DFS2(1);
T.Create(1, 1, n);
while(m--){
char Q[3];
int u, v, x;
scanf("%s%d%d", Q, &u, &v);
if(Q[0] == 'C') scanf("%d", &x), change(u, v, x);
else printf("%d\n", query(u, v));
}
return 0;
}