P2486 [SDOI2011]染色

#include<iostream>
#include<vector>
using namespace std;
#define int long long
const int N = 1e5 + 1;
int n , m ;
vector < int > to [ N ] ;
int dep [ N ] , fa [ N ] ,mson [ N ] , siz [ N ];
int value [ N ] ;
int top [ N ] ,vt [ N ], cnt, a [ N ] ;
void merge ( int x , int y ){ to [ x ] . push_back ( y ) ; to [ y ] . push_back ( x ) ; }
struct node {
    int lc , rc , color ;
    node(){lc = 0 , rc = 0 , color = 0 ;}
    node operator + ( const node & rhs ) const {
        node res ;
        if ( lc >  0 ) res . lc = lc;
        else
            res . lc = rhs . lc ;
        res . rc = rhs . rc;
        if ( rhs . lc == rc && rc > 0 ) res . color = -1;
        res . color += rhs . color;
        res . color += color;
        return res;
    }
} tree [ N << 2 ];
int tag [ N << 2 ];
namespace xds
{
    void build ( int k , int l , int r )
    {
        tag [ k ] = 0;
        if ( l == r ) {
                tree [ k ] . lc = tree [ k ] . rc = a [ l ] ;
                tree [ k ] . color = 1 ;
                return ;
        }
        int mid = l + r >> 1;
        build ( k << 1 , l , mid );
        build ( k << 1 | 1 , mid + 1 , r ) ;
        tree [ k ] = tree [ k << 1 ] + tree [ k << 1 | 1 ];
        return;
    }
    void push_down ( int k )
    {
        if ( tag [ k ] == 0 ) return;

        tree [ k << 1 ] . color = 1;
        tree [ k << 1 ] . lc = tree[ k << 1 ] . rc = tree [ k ] . lc;
        tag [ k << 1 ] = tag [ k ];

        tree [ k << 1 | 1 ] . color = 1 ;
        tree [ k << 1 | 1 ] . lc = tree [ k << 1 | 1 ] . rc = tree [ k ] . lc ;
        tag [ k << 1 | 1 ] = tag [ k ] ;

        tag [ k ] = 0;
    }
    void update ( int k , int l , int r , int x , int y , int c )
    {
        if ( l >= x && r <= y )
        {
            tree [ k ] . lc = tree [ k ] . rc = c;
            tree [ k ] . color = 1;
            tag [ k ] = c ;
            return;
        }
        int mid = l + r >> 1;
        push_down( k );
        if ( x <= mid ) update ( k << 1 , l , mid , x , y , c );
        if ( y > mid ) update ( k << 1 | 1 , mid + 1 , r , x , y , c ) ;
        tree [ k ] = tree [ k << 1 ] + tree [ k << 1 | 1 ];
        return ;
    }
    node query ( int k , int l , int r , int x , int y )
    {
        if ( l >= x && r <= y )
        {
            return tree [ k ]  ;
        }
        int mid = l + r >> 1;
        push_down ( k );
        node rhs ;
        if ( x <= mid ) rhs = rhs + query ( k << 1 , l , mid , x , y );
        if ( y > mid ) rhs = rhs + query ( k << 1 | 1 , mid + 1 , r , x , y );
        tree [ k ] = tree [ k << 1 ] + tree [ k << 1 | 1 ];
        return rhs ;
    }
}
namespace pf
{
    void dfs_1 ( int x , int father )
    {
        dep [ x ] = dep [ father ] + 1;
        fa [ x ] = father;
        siz [ x ] = 1;
        int maxsize = -1;
        for ( auto y : to [ x ] )
        {
            if ( y == father ) continue;
            dfs_1 ( y , x );
            siz [ x ] += siz [ y ] ;
            if ( siz [ y ] > maxsize )
                maxsize = siz [ y ] , mson [ x ] = y;
        }
        return ;
    }
    void dfs_2 ( int x , int topfa )
    {
        top [ x ] = topfa;
        vt [ x ] = ++ cnt;
        a [ cnt ] = value [ x ];
        if ( mson [ x ] > 0 )
            dfs_2 ( mson [ x ] , topfa );
        for ( auto y : to [ x ] )
        {
            if ( y == fa [ x ] ) continue;
            if ( y == mson [ x ] ) continue;
            dfs_2 ( y , y ) ;
        }
    }
    void cz_1 ( int x , int y , int z )
    {
        if ( dep [ x ] < dep [ y ] ) swap ( x , y );
        while ( top [ x ] != top [ y ] )
        {
            if ( dep [ top [ x ] ] <= dep [ top [ y ] ] ) swap ( x , y );
            xds :: update ( 1 , 1 , n , vt [ top [ x ] ] , vt [ x ] , z );
            x = top [ x ] ;
            x = fa [ x ] ;
        }
        if ( dep [ x ] < dep [ y ] ) swap ( x , y );
        xds :: update ( 1 , 1 , n , vt [ y ] , vt [ x ] , z ) ;
    }
    int cz_2 ( int x , int y )
    {
        node px , py ;
        while ( top [ x ] != top [ y ] ) {
            if ( dep [ top [ x ] ] <= dep [ top [ y ] ] ) {
                py = xds :: query ( 1 , 1 , n , vt [ top [ y ] ] , vt [ y ] ) + py;
                y = fa [ top [ y ] ];
            }
            else {
                px = xds :: query ( 1 , 1 , n , vt [ top [ x ] ] , vt [ x ] ) + px;
                x = fa [ top [ x ] ];
            }
        }
        if ( dep [ x ] > dep [ y ] )
        {
            px = xds :: query ( 1 , 1 , n , vt [ y ] , vt [ x ] ) + px;
            if ( px . lc == py . lc ) px . color -- ;
            return px . color + py . color;
        }
        else
        {
            py = xds :: query ( 1 , 1 , n , vt [ x ] , vt [ y ] ) + py;
            if ( px . lc == py . lc ) px . color -- ;
            return px . color + py . color;
        }
    }
}
signed main()
{
    freopen ( "text.in", "r" , stdin );
    cin >> n >> m;
    for ( int i = 1 ; i <= n ; i ++ )
        cin >> value [ i ] ;
    for ( int i = 1 ; i < n ; i ++ )
    {
        static int x , y ;
        cin >> x >> y;
        merge ( x , y );
    }
    pf :: dfs_1 ( 1 , 0 ) ;
    pf :: dfs_2 ( 1 , 1 );
    xds :: build ( 1 , 1 , n );
    for ( int i = 1 ; i <= m ; i++ )
    {
        char op ;
        int x , y , z;
        cin >> op >> x;
        if ( op == 'C' )
        {
            cin >> y >> z ;
            pf :: cz_1 ( x , y , z );
        }
        if ( op == 'Q' )
        {
            cin >> y;
            cout << pf :: cz_2 ( x , y ) << endl ;
        }
    }
    return 0;
}