洛谷P2486 [SDOI2011]染色（树链剖分+线段树判断边界）

【题目链接】

【思路】：

涉及到树上区间修改操作，所以使用树链剖分，涉及到区间查询，所以使用线段树。

update操作时，就正常操作，难点在于query操作的计数。

因为树链剖分的dfs序只能保证一条重链上的dfn[]连续，不能使得任意两点之间简单路径上的dfn[]连续，所以当x往上跳到fa[top[x]]时，要判断

top[x]的颜色与fa[top[x]]的颜色是否相同，如果相同要再减一。

以及在线段树中query操作和pushUp时，都要判断左儿子的右端点与右儿子的左端点是否相同，如果在pushUp中相同，则令此时的segtree[id]减一，

如果在query中相同，那么使得query函数返回的结果减一。 接下来上代码。

#include <bits/stdc++.h>
using namespace std;

const int maxn = 1e5 + 5;
int n, m, a[maxn];
int segtree[maxn<<2], lazy[maxn<<2], lft[maxn<<2], rht[maxn<<2];
int head[maxn], tot, cnt;
int fa[maxn], top[maxn], dfn[maxn], rk[maxn], siz[maxn], dep[maxn], son[maxn];
struct edge{
    int to,next;
} ed[maxn<<1];
inline int read(){
    int k=0, f=1;
    char ch=getchar();
    while( ch>'9' || ch<'0' ){ if(ch=='-') f=-1; ch=getchar(); }
    while( ch>='0' && ch<='9' ){ k=k*10+ch-'0'; ch=getchar(); }
    return k*f;
}

inline void init(){
    memset( head ,-1, sizeof(head) );
    memset( lazy, -1, sizeof(lazy) );
    tot = 1;
}

inline void add( int u, int v ){
    ed[++tot].to = v; ed[tot].next = head[u]; head[u] = tot;
    ed[++tot].to = u; ed[tot].next = head[v]; head[v] = tot;
}

inline void dfs1( int x ){
    siz[x] = 1;
    for( int i=head[x]; ~i; i=ed[i].next ){
        int y = ed[i].to;
        if( y==fa[x] ) continue;
        dep[y] = dep[x]+1;
        fa[y] = x;
        dfs1(y);
        siz[x] += siz[y];
        if( son[x]==0 || siz[y]>siz[son[x]] ) son[x] = y;
    }
}

inline void dfs2( int x, int tp ){
    top[x] = tp;
    dfn[x] = ++cnt;
    rk[cnt] = x;
    if( son[x] ) dfs2( son[x], tp );
    for( int i=head[x]; ~i; i=ed[i].next ){
        int y = ed[i].to;
        if( y!=fa[x] && y!=son[x] ) dfs2(y, y);
    }
}

inline void pushUp( int id ){
    segtree[id] = segtree[id<<1] + segtree[id<<1|1];
    if( lft[id<<1|1]==rht[id<<1] ) segtree[id] --;
    lft[id] = lft[id<<1]; rht[id] = rht[id<<1|1];
}

inline void pushDown( int id ){
    if( lazy[id]==-1 ) return;
    lazy[id<<1] = lazy[id<<1|1] = lazy[id];
    segtree[id<<1] = segtree[id<<1|1] = 1;
    lft[id<<1] = lft[id<<1|1] = rht[id<<1] = rht[id<<1|1] = lazy[id];
    lazy[id] = -1;
}

inline void build( int l, int r, int id ){
    if( l==r ){
        lft[id] = rht[id] = a[rk[l]];
        segtree[id] = 1;
        return ;
    }
    int mid=l+r>>1;
    build( l, mid, id<<1 );
    build( mid+1, r, id<<1|1 );
    pushUp(id);
}

inline void update_tree( int l, int r, int ql, int qr, int id, int c ){
    if( ql<=l && qr>=r ){
        segtree[id] = 1;
        lft[id] = rht[id] = c;
        lazy[id] = c;
        return ;
    }
    pushDown(id);
    int mid = l+r>>1;
    if( ql<=mid ) update_tree( l, mid, ql, qr, id<<1, c );
    if( qr>mid ) update_tree( mid+1, r, ql, qr, id<<1|1, c );
    pushUp(id);
}

inline int query( int l, int r, int ql, int qr, int id ){
    if( ql<=l && qr>=r ) return segtree[id];
    pushDown(id);
    int mid = l+r>>1;
    int res = 0;
    if( ql<=mid ) res += query( l, mid, ql, qr, id<<1 );
    if( qr>mid ) res += query( mid+1, r, ql, qr, id<<1|1 );
    if( ql<=mid && qr>mid && lft[id<<1|1]==rht[id<<1] ) res--;      //这里也要判断一次是否相同,前提是要跨立区间
    return res;
}

inline int color_query( int l, int r, int idx, int id ){
    if( l==r ) return lft[id];
    pushDown(id);
    int mid = l+r>>1;
    if( idx<=mid ) return color_query( l, mid, idx, id<<1 );
    else return color_query( mid+1, r, idx, id<<1|1 );
}

inline void swap( int &x, int &y ){ x^=y^=x^=y; }

inline int sum( int x, int y ){
    int res = 0;
    while( top[x]!=top[y] ){
        if( dep[top[x]]<dep[top[y]] ) swap(x, y);
        res += query(1, n, dfn[top[x]], dfn[x], 1 );
        if( color_query(1, n, dfn[top[x]], 1)==color_query(1, n, dfn[fa[top[x]]], 1) ) res --;  //颜色相同减一
        x = fa[top[x]];
    }
    if( dfn[x]>dfn[y] ) swap(x, y);
    res += query( 1, n, dfn[x], dfn[y], 1 );
    return res==0 ? 1:res;
}

inline void update_chain( int x, int y, int c ){
    while( top[x]!=top[y] ){
        if( dep[top[x]]<dep[top[y]] ) swap(x, y);
        update_tree(1, n, dfn[top[x]], dfn[x], 1, c );
        x = fa[top[x]];
    }
    if( dfn[x]>dfn[y] ) swap(x, y);
    update_tree( 1, n, dfn[x], dfn[y], 1, c );
}

int main(){
    // freopen("in.txt", "r", stdin);
    init();
    n = read(); m = read();
    for( int i=1; i<=n; i++ ) a[i] = read();
    for( int i=1; i<n; i++ ){
        int u=read(), v=read();
        add(u, v);
    }
    dep[1] = fa[1] = 1;
    dfs1(1);
    dfs2(1, 1);
    build( 1, n, 1 );
    char ch[5];
    while( m-- ){
        int x, y, z;
        scanf("%s", ch);
        x = read(); y = read();
        if( ch[0]=='Q' ) printf("%d\n", sum(x, y));
        else{
            z = read();
            update_chain(x, y, z);
        }
    }

    return 0;
}