[SDOI2011]染色

题目描述

输入格式

不只有线段树外的节点需要合并，线段树内相邻节点，也要判合并

#include<bits/stdc++.h>
#define re return
#define inc(i,l,r) for(int i=l;i<=r;++i)

using namespace std;
template<typename T>inline void rd(T&x)
{
    char c;bool f=0;
    while((c=getchar())<'0'||c>'9')if(c=='-')f=1;
    x=c^48;
    while((c=getchar())>='0'&&c<='9')x=x*10+(c^48);
    if(f)x=-x;
}

const int maxn=100005; 
int n,m,tot,ans,k=1;
int size[maxn],son[maxn],top[maxn],fa[maxn],deep[maxn];
int rev[maxn],seg[maxn],col[maxn];
int hd[maxn];

struct node{
    int to,nt;
}e[maxn<<1]; 
inline void add(int x,int y)
{
    e[++k].to=y;e[k].nt=hd[x];hd[x]=k; 
    e[++k].to=x;e[k].nt=hd[y];hd[y]=k; 
}

//-------------------------------------------------------------------------------
inline void dfs1(int x)
{
    size[x]=1;
    deep[x]=deep[fa[x]]+1;
    for(int i=hd[x];i;i=e[i].nt)
    {
        int v=e[i].to;
        if(v==fa[x])continue;
        fa[v]=x;
        dfs1(v);
        size[x]+=size[v];
        if(size[v]>size[son[x]])
            son[x]=v;
    }
    
}

inline void dfs2(int x,int ftop)
{
    top[x]=ftop;
    seg[x]=++tot;
    rev[tot]=x;
    if(son[x])
        dfs2(son[x],ftop);
    for(int i=hd[x];i;i=e[i].nt)
    {
        int v=e[i].to;
        if(!top[v])
            dfs2(v,v);
    }
}
//----------------------------------------------------------------------------------
#define lson rt<<1
#define rson rt<<1|1
int sum[maxn<<2],lc[maxn<<2],rc[maxn<<2],lazy[maxn<<2];

inline void pushup(int rt)
{
    lc[rt]=lc[lson];
    rc[rt]=rc[rson];
    sum[rt]=sum[lson]+sum[rson]-(rc[lson]==lc[rson]);
}

inline void pushdown(int rt)
{
    lazy[lson]=lazy[rson]=1;
    sum[lson]=1;sum[rson]=1;
    lc[lson]=rc[lson]=lc[rt];
    lc[rson]=rc[rson]=rc[rt];
    lazy[rt]=0;
}

inline void build(int rt,int l,int r)
{
    if(l==r)
    {
        lc[rt]=rc[rt]=col[rev[l]];
        sum[rt]=1;
        re ;
    }
    int mid=(l+r)>>1;
    build(lson,l,mid);
    build(rson,mid+1,r);
    pushup(rt);
}

inline void addd(int rt,int l,int r,int x,int y,int z)
{
    if(x<=l&&r<=y)
    {
        sum[rt]=1;
        lazy[rt]=1;
        lc[rt]=rc[rt]=z;
        re ;
    }
    if(lazy[rt])pushdown(rt);
    int mid=(l+r)>>1;
    if(x<=mid)addd(lson,l,mid,x,y,z);
    if(y>mid)addd(rson,mid+1,r,x,y,z);
    pushup(rt);
}


struct qwq{
    int l_col,r_col;
};
inline qwq query(int rt,int l,int r,int x,int y)
{
    if(x<=l&&r<=y)
    {
        ans+=sum[rt];
        re (qwq){lc[rt],rc[rt]};
    }
    int mid=(l+r)>>1;
    if(lazy[rt])pushdown(rt);
    if(y<=mid)re query(lson,l,mid,x,y);
    if(mid<x) re query(rson,mid+1,r,x,y);
    else 
    {
        qwq v1=query(lson,l,mid,x,y);
        qwq v2=query(rson,mid+1,r,x,y);
        
        if(v1.r_col==v2.l_col)--ans;
        
        re (qwq){v1.l_col,v2.r_col};
    }
}


int main()
{
//    freopen("in.txt","r",stdin);
    
    int x,y,z;
    rd(n),rd(m);
    inc(i,1,n)
    rd(col[i]);
    
    inc(i,2,n)
    {
        rd(x),rd(y);
        add(x,y);
    }
    
    dfs1(1);
    dfs2(1,1);
    build(1,1,n);
    
    
    int cnt=0;
    char ss[10];
    inc(i,1,m)
    {
        scanf("%s",ss);
        rd(x),rd(y);
        if(ss[0]=='Q')
        {
            ++cnt;
            ans=0;
        
            int lastx=-1,lasty=-1;    
              while(top[x]!=top[y])
              {
                  if(deep[top[x]]<deep[top[y]])
                  {
                      swap(x,y);
                      swap(lastx,lasty);
                }
                
            qwq    v1=query(1,1,n,seg[top[x]],seg[x]);
                ans-=(lastx==v1.r_col);
                lastx=v1.l_col;
                x=fa[top[x]];
            }
            
            if(deep[x]>deep[y])
            {
                swap(x,y);
                swap(lastx,lasty);
            }
            
            qwq v1=query(1,1,n,seg[x],seg[y]);
            ans=ans-(lastx==v1.l_col)-(lasty==v1.r_col);
            printf("%d\n",ans);
        }
        else
        {
             rd(z);
             while(top[x]!=top[y])
             {
                 if(deep[top[x]]<deep[top[y]])
                    swap(x,y);
                addd(1,1,n,seg[top[x]],seg[x],z);
                x=fa[top[x]];
            }
            if(deep[x]>deep[y])swap(x,y);
            addd(1,1,n,seg[x],seg[y],z);
        }
    }
    re 0;
}