POJ 2243 [SDOI2011]染色 | 树链剖分+线段树

原题链接

---

 

肯定是树链剖分的题啦

树剖怎么做可以看我上一篇博客

如果我们已经剖完了:

然后考虑怎么维护重链和查询

用线段树维护的时候当前区间的区间颜色个数应该等于左儿子+右儿子,但是当左儿子的右端点和右儿子的左端点颜色一样,显然区间数要减1

所以每个节点存一下左端点的右端点颜色正常维护即可

考虑查询,在同一重链上的点显然线段树可以解决,当top[u]!=top[v]的时候,得让deep较深的爬树,那么答案就要+=爬树那一段的区间个数

但是有可能的是fa[top[u]]的颜色和top[u]的颜色相等,这个时候需要答案--,这样才能保证下次爬树的时候不会多算

注意:请写单点查询询问颜色,因为fa[top[u]]和top[u]是轻链的两端,不能保证在线段树中编号连续,所以不能区间查询是否是个数是1

修改同理,爬树的时候修改即可

#include<cstdio>
#include<algorithm>
#include<cstring>
#define N 100010
using namespace std;
int n,m,head[N],indx[N],pos[N],fa[N],ecnt,color[N],deep[N],sz[N],top[N],a,b,c,tot,son[N];
char s[N];
int read()
{
    int ret=0,neg=1;
    char j=getchar();
    for (;j>'9' || j<'0';j=getchar())
    if (j == '-') neg=-1;
    for (;j>='0' && j<='9';j=getchar())
    ret=ret*10+j-'0';
    return ret*neg;
}
struct adj
{
    int nxt,v;
}e[2*N];
struct node
{
    int l,r,Lcolor,Rcolor,sum,lz;
}t[4*N];
void add(int u,int v)//加边
{
    e[++ecnt].v=v;
    e[ecnt].nxt=head[u];
    head[u]=ecnt;
    e[++ecnt].v=u;
    e[ecnt].nxt=head[v];
    head[v]=ecnt;
}
void dfs1(int x,int father,int d)//第一次dfs
{
    fa[x]=father,deep[x]=d,sz[x]=1;
    for (int i=head[x];i;i=e[i].nxt)
    {
    int v=e[i].v;
    if (v==father) continue;
    dfs1(v,x,d+1);
    sz[x]+=sz[v];
    if (sz[son[x]]<sz[v]) son[x]=v;
    }
}
void dfs2(int x,int TOP)//第二次dfs
{
    pos[x]=++tot;
    indx[tot]=x;
    top[x]=TOP;
    if (son[x]) dfs2(son[x],TOP);
    for (int i=head[x];i;i=e[i].nxt)
    {
    int v=e[i].v;
    if (v==fa[x] || v==son[x]) continue;
    if (v!=0) dfs2(v,v);
    }
}
void pushup(int p)//线段树更新
{
    t[p].sum=t[p<<1].sum+t[p<<1|1].sum;
    if (t[p<<1].Rcolor==t[p<<1|1].Lcolor) t[p].sum--;
    t[p].Lcolor=t[p<<1].Lcolor;
    t[p].Rcolor=t[p<<1|1].Rcolor;
}
void pushdown(int p)//lazy下放
{
    if (t[p].l==t[p].r || t[p].lz==-1) return;
    int w=t[p].lz;
    t[p<<1].Lcolor=t[p<<1].Rcolor=t[p<<1|1].Lcolor=t[p<<1|1].Rcolor=t[p<<1].lz=t[p<<1|1].lz=w;
    t[p<<1].sum=t[p<<1|1].sum=1;
    t[p].lz=-1;
}
void build(int p,int l,int r)//建树
{
    t[p].l=l,t[p].r=r,t[p].lz=-1;
    if (l==r)
    {
    t[p].Lcolor=t[p].Rcolor=color[indx[l]];
    t[p].sum=1;
    }
    else
    {
    int mid=l+r>>1;
    build(p<<1,l,mid);
    build(p<<1|1,mid+1,r);
    pushup(p);
    }
}
void modify(int p,int l,int r,int k)//区间修改
{
    if (l==t[p].l && r==t[p].r)
    {
    t[p].sum=1;
    t[p].Lcolor=t[p].Rcolor=k;
    t[p].lz=k;
    return;
    }
    pushdown(p);
    int mid=t[p].l+t[p].r>>1;
    if (r<=mid)
    modify(p<<1,l,r,k);
    else if (l>mid) modify(p<<1|1,l,r,k);
    else
    modify(p<<1,l,mid,k),modify(p<<1|1,mid+1,r,k);
    pushup(p);
}
int query(int p,int l,int r)//区间询问
{
    if (l==t[p].l && r==t[p].r)
    return t[p].sum;
    pushdown(p);
    int mid=t[p].l+t[p].r>>1;
    if (r<=mid)
    return query(p<<1,l,r);
    if (l>mid) return query(p<<1|1,l,r);
    int tmp=query(p<<1,l,mid)+query(p<<1|1,mid+1,r);
    if (t[p<<1].Rcolor==t[p<<1|1].Lcolor) tmp--;
    return tmp;
}
int Qcolor(int p,int l)//单点询问
{
    if (t[p].l==t[p].r && t[p].l==l) return t[p].Lcolor;
    pushdown(p);
    int mid=t[p].l+t[p].r>>1;
    if (l<=mid) return Qcolor(p<<1,l);
    else return Qcolor(p<<1|1,l);
}
void getcolor(int a,int b,int w)//修改
{
    while (top[a]!=top[b])//爬树
    {
    if (deep[top[a]]<deep[top[b]]) swap(a,b);
    modify(1,pos[top[a]],pos[a],w);
    a=fa[top[a]];
    }
    if (deep[a]>deep[b]) swap(a,b);
        modify(1,pos[a],pos[b],w);
}
int querycolor(int u,int v)//询问
{
    int ret=0;
    while (top[u]!=top[v])
    {
    if (deep[top[u]]<deep[top[v]]) swap(u,v);
    ret+=query(1,pos[top[u]],pos[u]);
    if (Qcolor(1,pos[fa[top[u]]])==Qcolor(1,pos[top[u]])) ret--;//看题解
    u=fa[top[u]];
    }
    if (deep[u]>deep[v]) swap(u,v);
    return ret+query(1,pos[u],pos[v]);
}
int main()
{
    n=read(),m=read();
    for (int i=1;i<=n;i++)
    color[i]=read();
    for (int i=1;i<n;i++)
    add(read(),read());
    dfs1(1,0,0);
    dfs2(1,1);
    build(1,1,n);
    while (m--)
    {
    scanf("%s",s);
    if (s[0]=='C')
    {
        a=read(),b=read(),c=read();
        getcolor(a,b,c);
    }
    else
    {
        scanf("%d%d",&a,&b);
        printf("%d\n",querycolor(a,b));
    }
    }
    return 0;
}