BZOJ2243: [SDOI2011]染色

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]之间。

Solution

很有意思和恶心的一道题。线段树+树剖维护颜色段。
挺直观的一个想法就是维护sum和lcol和rcol，然后合并的时候左儿子的rcol如果等于右儿子的lcol的话就-1（成一段了）。
然后在树上面跑的时候注意维护之前的最左边颜色，然后在query的过程中顺便存下来当前这段的lcol和rcol（这个开两个全局变量就好了），然后考虑按树剖的dfs序编号时是怎么样的：重链先。所以每次记录左右两个端点往上跳时的最左端点颜色，swap的时候一起swap就好了。
然后注意跳到lca之后，因为深度小的那个点一定是lca，而且两者此时在同一条重链上，所以要判断重合段又会比较麻烦，具体是：如果当前query的这一段的右端点等于lca那个点上次查询的右端点就--ans如果，如果当前这段的query的左端点等于深度大的那段的右端点，就要再--ans
细节很多

#include <bits/stdc++.h>

#define ll long long
#define inf 0x3f3f3f3f 
#define il inline 

namespace io {

    #define in(a) a=read()
    #define out(a) write(a)
    #define outn(a) out(a),putchar('
')

    #define I_int int 
    inline I_int read() {
        I_int x = 0 , f = 1 ; char c = getchar() ;
        while( c < '0' || c > '9' ) { if( c == '-' ) f = -1 ; c = getchar() ; } 
        while( c >= '0' && c <= '9' ) { x = x * 10 + c - '0' ; c = getchar() ; } 
        return x * f ;
    } 
    char F[ 200 ] ;
    inline void write( I_int x ) {
        if( x == 0 ) { putchar( '0' ) ; return ; }
        I_int tmp = x > 0 ? x : -x ;
        if( x < 0 ) putchar( '-' ) ;
        int cnt = 0 ;
        while( tmp > 0 ) {
            F[ cnt ++ ] = tmp % 10 + '0' ;
            tmp /= 10 ;
        }
        while( cnt > 0 ) putchar( F[ -- cnt ] ) ;
    }
    #undef I_int

}
using namespace io ;

using namespace std ;

#define N 100010

int cnt, head[N];
struct edge {
    int to, nxt;
}e[N<<1];
int dep[N], fa[N], top[N], a[N], id[N], siz[N], w[N];
int n, m;

void ins(int u, int v) {
    e[++cnt] = (edge) {v, head[u]};
    head[u] = cnt;
}

void dfs1(int u) {
    siz[u] = 1;
    for(int i = head[u]; i; i = e[i].nxt) {
        int v = e[i].to;
        if(v == fa[u]) continue;
        dep[v] = dep[u] + 1;
        fa[v] = u;
        dfs1(v);
        siz[u] += siz[v];
    }
}

int tim;
void dfs2(int u, int topf) {
    top[u] = topf;
    id[u] = ++tim;
    w[tim] = a[u];
    int k = 0;
    for(int i = head[u]; i; i = e[i].nxt) {
        if(e[i].to != fa[u] && siz[e[i].to] > siz[k]) k = e[i].to; 
    }
    if(!k) return;
    dfs2(k, topf);
    for(int i = head[u]; i; i = e[i].nxt) {
        if(e[i].to != fa[u] && e[i].to != k) dfs2(e[i].to, e[i].to);
    }
}

// seg-tree
struct tree {
    int l, r, lcol, rcol;
    int v, tag;
}t[N<<2];
#define lc (rt << 1)
#define rc (rt << 1 | 1)

void up(int rt) {
    t[rt].lcol = t[lc].lcol; t[rt].rcol = t[rc].rcol;
    int num = t[lc].v + t[rc].v;
    if(t[lc].rcol == t[rc].lcol) --num;
    t[rt].v = num;
}

void build(int l, int r, int rt) {
    t[rt].l = l; t[rt].r = r;
    if(l == r) {
        t[rt].lcol = t[rt].rcol = w[l];
        t[rt].v = 1;
        return;
    }
    int mid = (l + r) >> 1;
    build(l, mid, lc);
    build(mid + 1, r, rc);
    up(rt);
}
#define l (t[rt].l)
#define r (t[rt].r)
#define mid ((l + r) >> 1)

void down(int rt) {
    if(t[rt].tag) {
        t[lc].tag = t[rc].tag = t[rt].tag;
        t[lc].v = t[rc].v = 1;
        t[lc].lcol = t[rc].lcol = t[rt].lcol;
        t[lc].rcol = t[rc].rcol = t[rt].rcol;
        t[rt].tag = 0;
    }
}

void upd(int L, int R, int c, int rt) {
    if(L <= l && r <= R) {
        t[rt].v = 1;
        t[rt].lcol = t[rt].rcol = c;
        t[rt].tag = c;
        return;
    }
    down(rt);
    if(L <= mid) upd(L, R, c, lc);
    if(R > mid) upd(L, R, c, rc);
    up(rt);
}

int now_l, now_r;
int query(int L, int R, int rt) {
    if(L == l) now_l = t[rt].lcol;
    if(R == r) now_r = t[rt].rcol;
    if(L <= l && r <= R) return t[rt].v;
    int ans = 0, tot = 0;
    down(rt);
    if(L <= mid) ans += query(L, R, lc), tot++;
    if(R > mid) ans += query(L, R, rc), tot++;
    if(tot == 2) if(t[lc].rcol == t[rc].lcol) --ans;
    return ans;
}

#undef mid
#undef lc
#undef rc
#undef l
#undef r
// seg-tree end

// tree-chain 
void upd_chain(int x, int y, int c) {
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        upd(id[top[x]], id[x], c, 1);
        x = fa[top[x]];
    }
    if(dep[x] > dep[y]) swap(x, y);
    upd(id[x], id[y], c, 1);
}

int query_chain(int x, int y) {
    int col_l = -1, col_r = -1, ans = 0;
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x, y), swap(col_l, col_r);
        ans += query(id[top[x]], id[x], 1);
        if(now_r == col_l) --ans;
        col_l = now_l;
        x = fa[top[x]];
    }
    if(dep[x] > dep[y]) swap(x, y), swap(col_l, col_r);
    ans += query(id[x], id[y], 1);
    if(now_r == col_r) --ans;
    if(now_l == col_l) --ans;
    return ans;
}
// tree-chain end

int main() {
    in(n), in(m);
    for(int i = 1; i <= n; ++i) in(a[i]);
    for(int i = 1; i < n; ++i) {
        int u = read(), v = read();
        ins(u, v), ins(v, u);
    }
    dfs1(1), dfs2(1, 1); 
    build(1, n, 1);
    for(int i = 1; i <= m; ++i) {
        char ch[2]; scanf("%s", ch);
        int x, y, v; in(x), in(y);
        if(ch[0] == 'C') {
            in(v);
            upd_chain(x, y, v);
        } else outn(query_chain(x, y));
    }
    return 0;
}