BZOJ2243 【SDOI2011】染色 树链剖分

Description


给定一棵有n个节点的无根树和m个操作，操作有2类：
1、将节点x到节点y路径上所有点染色成颜色c
2、询问节点x到节点y路径上的颜色段数量


Solution


先考虑在序列上求颜色段数量的问题。可以在线段树上维护每个区间左右端点的颜色和区间内颜色段的数量，然后有
：

tree[p].sum = tree[p<<1].sum + tree[p<<1|1].sum - ( tree[p<<1].rc == tree[p<<1|1].lc);

那么到树上后就再套一个树链剖分就好了


Code:

#include
#include
#include
#include
using namespace std;
const int maxn = 100005;
const int INF  = 0x3f3f3f3f;
struct edge{
    int x,y,next;
    edge(){}
    edge(int _x,int _y,int _nt):x(_x),y(_y),next(_nt){}
} e[maxn << 1];
int head[maxn] , tot = 0;
int siz[maxn],son[maxn],fa[maxn],dep[maxn],p[maxn],top[maxn];
int a[maxn],b[maxn],n,T,m;
inline void addedge(int x,int y){
    e[++tot] = edge( x , y , head[x] ); head[x] = tot;
    e[++tot] = edge( y , x , head[y] ); head[y] = tot;
}
void DFS1(int x){
    siz[x] = 1; son[x] = 0;
    for(int y,i = head[x] ; i ; i = e[i].next){
        y = e[i].y;
        if( dep[y] ) continue;
        fa[y] = x; dep[y] = dep[x] + 1;
        DFS1( y );
        siz[x] += siz[y];
        if( siz[y] > siz[son[x]] ) son[x] = y;
    }
}
void DFS2(int x,int chain){
    p[x] = ++T; top[x] = chain;
    if( son[x] ) DFS2( son[x] , chain );
    for(int y,i = head[x] ; i ; i = e[i].next ){
        y = e[i].y;
        if( p[y] || y == son[x] ) continue;
        DFS2( y , y );
    }
}

struct node{
    int L,R;
    int lc,rc,sum;
    int fg;
}tree[maxn << 2];
void update(int p,int c){
    tree[p].fg = c;
    tree[p].lc = tree[p].rc = c;
    tree[p].sum = 1;
}
void pushup(int p){
    tree[p].lc = tree[p<<1].lc;
    tree[p].rc = tree[p<<1|1].rc;
    tree[p].sum = tree[p<<1].sum + tree[p<<1|1].sum;
    if( tree[p<<1].rc == tree[p<<1|1].lc ) tree[p].sum--;
}
void pushdown(int p){
    int fg = tree[p].fg;
    if(fg == 0)return;
    tree[p].fg = 0;
    update( p<<1 , fg );
    update( p<<1|1 , fg );
}
void build(int p,int L,int R){
    tree[p].L = L; tree[p].R = R;
    tree[p].fg = 0;
    if( L == R ){
        tree[p].lc = tree[p].rc = a[L];
        tree[p].sum = 1;
        return;
    }
    int mid = ( L + R ) >> 1;
    build( p<<1 , L , mid );
    build( p<<1|1 , mid + 1 , R );
    pushup(p);
}
void change(int p,int L,int R,int c){
    int n_L = tree[p].L, n_R = tree[p].R;
    if( L <= n_L && n_R <= R){
        update( p , c );
        return;
    }
    pushdown( p );
    int mid = ( n_L + n_R ) >> 1;
    if( L <= mid ) change( p<<1 , L , R , c );
    if( mid < R ) change( p<<1|1 , L , R , c );
    pushup( p );
}
int query_sum(int p,int L,int R){
    int n_L = tree[p].L , n_R = tree[p].R;
    if( L <= n_L && n_R <= R ) return tree[p].sum;
    pushdown( p );
    int mid = ( n_L + n_R ) >> 1;
    int sum = 0;
    if( L <= mid ) sum += query_sum( p<<1 , L , R ); 
    if( mid < R ) sum += query_sum( p<<1|1 , L , R );
    if( ( L <= mid ) && ( mid < R ) && ( tree[p<<1].rc == tree[p<<1|1].lc) ) sum --;
    return sum;
}
int get_c(int p,int pos){
    int L = tree[p].L , R = tree[p].R;
    if(L == R) return tree[p].lc;
    pushdown( p );
    int mid = ( L + R ) >> 1;
    if( pos <= mid ) return get_c( p<<1 , pos );
    else return get_c( p<<1|1 ,pos );
}
void solve_change(int x,int y,int c){
    while( top[x] != top[y] ){
        if( dep[top[x]] < dep[top[y]] ) swap(x , y );
        change( 1 , p[top[x]] , p[x] , c );
        x = fa[top[x]];
    }
    if( dep[x] > dep[y] ) swap( x , y );
    change( 1 , p[x] , p[y] , c );
}
void solve_query(int x,int y){
    int sum = 0;
    while( top[x] != top[y] ){
        if( dep[top[x]] < dep[top[y]] ) swap( x , y );
        sum += query_sum( 1 , p[top[x]] , p[x] );
        if( get_c( 1 , p[top[x]] ) == get_c( 1 , p[fa[top[x]]] ) ) sum--;
        x = fa[top[x]];
    }
    if( dep[x] > dep[y] ) swap( x , y );
    sum += query_sum( 1 , p[x] , p[y] );
    printf("%d\n",sum);
}

int main(){
    scanf("%d%d",&n,&m);
    for(int i = 1 ; i <= n ; i++ ){
        scanf("%d",&b[i]);b[i]++;
    }

    for(int x,y,i = 1 ; i <= n-1 ; i++){
        scanf("%d%d",&x,&y);
        addedge(x,y);
    }
    dep[1] = 1;
    DFS1( 1 );
    DFS2( 1 , 1 );

    for(int i = 1 ; i <= n ; i++)
        a[p[i]] = b[i];

    build( 1 , 1 , n );

    char str[2];
    int x,y,c;
    for(int i = 1 ; i <= m ; i++){
        scanf("%s",str);
        if( str[0] == 'Q' ){
            scanf("%d%d",&x,&y);
            solve_query( x , y );
        }
        else if( str[0] == 'C' ){
            scanf("%d%d%d",&x,&y,&c);c++;
            solve_change( x , y , c );
        }
    }

}