洛谷 P2486 [SDOI2011]染色（线段树区间更新，树链剖分）

链接： 洛谷 P2486 [SDOI2011]染色

题目描述

输入格式

输出格式
对于每个询问操作，输出一行答案。

输入输出样例
输入 #1

6 5
2 2 1 2 1 1
1 2
1 3
2 4
2 5
2 6
Q 3 5
C 2 1 1
Q 3 5
C 5 1 2
Q 3 5

输出 #1

3
1
2

---

分析：

树上区间操作，首先进行树链剖分。

对于每一条链，线段树进行区间维护，对于 $[l,r]$（树剖后的新编号）的 颜色段数量 $sum$

• 若 左区间右端颜色 和 右区间左端颜色 不同，则有： $sum[l,r]=sum[l,mid]+sum[mid+1,r]$
• 若 左区间右端颜色 和 右区间左端颜色 相同，则有： $sum[l,r]=sum[l,mid]+sum[mid+1,r]-1$

所以，线段树还需要维护每个区间的 左端颜色 $L$ 和 右端颜色 $R$。

而对于查询到的不同重链的颜色段之和，同样地，要 判断连接处的颜色是否相同，即 判断一条重链链首和其父结点的颜色是否相同，若相同，则令求和的答案 $-1$。

---

以下代码：

#include<bits/stdc++.h>
#define LL long long
using namespace std;
const int INF=0x3f3f3f3f;
const int maxn=1e5+50;
int n,q,w[maxn];
//**************建树**************************
struct edge
{
    int u;
    int v;
    int next;
}e[maxn<<1];
int head[maxn],cnt;
void add_edge(int u,int v)
{
    e[cnt]=edge{u,v,head[u]};
    head[u]=cnt++;
    e[cnt]=edge{v,u,head[v]};
    head[v]=cnt++;
}
//**************树链剖分***********************
int fa[maxn],dep[maxn],sz[maxn],son[maxn];
void dfs1(int u,int pre,int depth)
{
    fa[u]=pre;
    dep[u]=depth;
    sz[u]=1;
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v=e[i].v;
        if(v==pre)
            continue;
        dfs1(v,u,depth+1);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]])
            son[u]=v;
    }
}
int id[maxn],top[maxn],wt[maxn],tot;
void dfs2(int u,int TOP)
{
    top[u]=TOP;
    id[u]=++tot;
    wt[id[u]]=w[u];
    if(!son[u])
        return;
    dfs2(son[u],TOP);
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v=e[i].v;
        if(v!=son[u]&&v!=fa[u])
            dfs2(v,v);
    }
}
//******************线段树*********************
#define ls rt<<1
#define rs rt<<1|1
struct seg_tree
{
    int sum;   //l~r共多少颜色段
    int laz;   //待向下传递的填入颜色
    int L;     //l~r中最左端的颜色
    int R;     //l~r中最右端的颜色
}t[maxn<<2];
void push_up(int rt)     //向上传递
{
    t[rt].sum=t[ls].sum+t[rs].sum;
    if(t[ls].R==t[rs].L)
        t[rt].sum--;
    t[rt].L=t[ls].L;
    t[rt].R=t[rs].R;
}
void push_down(int rt)   //懒标记laz向下传递
{
    if(!t[rt].laz)
        return;
    t[ls].sum=1;
    t[ls].laz=t[rt].laz;
    t[ls].L=t[ls].R=t[rt].laz;
    t[rs].sum=1;
    t[rs].laz=t[rt].laz;
    t[rs].L=t[rs].R=t[rt].laz;
    t[rt].laz=0;
}
void build(int rt,int l,int r)
{
    t[rt].laz=0;   //懒标记清空
    if(l==r)
    {
        t[rt].sum=1;
        t[rt].L=t[rt].R=wt[l];
        return;
    }
    int mid=(l+r)>>1;
    build(ls,l,mid);
    build(rs,mid+1,r);
    push_up(rt);
}
void updata(int rt,int l,int r,int ql,int qr,int val)
{
    if(ql<=l&&r<=qr)
    {
        t[rt].sum=1;    //l~r全涂为一种颜色，颜色段数量为1
        t[rt].L=t[rt].R=val;
        t[rt].laz=val;
        return;
    }
    push_down(rt);
    int mid=(l+r)>>1;
    if(ql<=mid)
        updata(ls,l,mid,ql,qr,val);
    if(qr>mid)
        updata(rs,mid+1,r,ql,qr,val);
    push_up(rt);
}
int query(int rt,int l,int r,int ql,int qr)
{
    if(ql<=l&&r<=qr)
        return t[rt].sum;
    push_down(rt);
    int mid=(l+r)>>1;
    if(qr<=mid)
        return query(ls,l,mid,ql,qr);
    else if(ql>mid)
        return query(rs,mid+1,r,ql,qr);
    else
    {
        int res=query(ls,l,mid,ql,qr)+query(rs,mid+1,r,ql,qr);
        if(t[ls].R==t[rs].L)   //合并时也要判断连接处颜色是否相同
            res--;
        return res;
    }
}
int query_color(int rt,int l,int r,int p)   //单点询问颜色
{
    if(l==r)
        return t[rt].L;
    push_down(rt);
    int mid=(l+r)>>1;
    if(p<=mid)
        return query_color(ls,l,mid,p);
    else
        return query_color(rs,mid+1,r,p);
}
//******************树上区间操作******************
void tree_updata(int x,int y,int val)   //从x结点到y结点全部更新
{
    while(top[x]!=top[y])
    {
        if(dep[top[x]]>=dep[top[y]])
        {
            updata(1,1,n,id[top[x]],id[x],val);
            x=fa[top[x]];
        }
        else
        {
            updata(1,1,n,id[top[y]],id[y],val);
            y=fa[top[y]];
        }
    }
    if(id[x]<=id[y])
        updata(1,1,n,id[x],id[y],val);
    else
        updata(1,1,n,id[y],id[x],val);
}
int tree_query(int x,int y)
{
    int res=0;
    while(top[x]!=top[y])
    {
        if(dep[top[x]]>=dep[top[y]])
        {
            res+=query(1,1,n,id[top[x]],id[x]);
            int t1=query_color(1,1,n,id[top[x]]),t2=query_color(1,1,n,id[fa[top[x]]]);
            if(t1==t2)       //判断两重链连接处颜色是否相同
                res--;
            x=fa[top[x]];
        }
        else
        {
            res+=query(1,1,n,id[top[y]],id[y]);
            int t1=query_color(1,1,n,id[top[y]]),t2=query_color(1,1,n,id[fa[top[y]]]);
            if(t1==t2)
                res--;
            y=fa[top[y]];
        }
    }
    if(id[x]<=id[y])
        res+=query(1,1,n,id[x],id[y]);
    else
        res+=query(1,1,n,id[y],id[x]);
    return res;
}
//*****************初始化**********************
void init()
{
    memset(head,-1,sizeof(head));
    cnt=0;
    memset(son,0,sizeof(son));
    tot=0;
}
//******************************************
int main()
{
    init();
    scanf("%d %d",&n,&q);
    for(int i=1;i<=n;i++)
        scanf("%d",&w[i]);
    for(int i=1;i<=n-1;i++)
    {
        int a,b;
        scanf("%d %d",&a,&b);
        add_edge(a,b);
    }
    dfs1(1,-1,1);
    dfs2(1,1);
    build(1,1,n);
    while(q--)
    {
        char op[5];
        int a,b,c;
        scanf("%s",op);
        if(op[0]=='C')
        {
            scanf("%d %d %d",&a,&b,&c);
            tree_updata(a,b,c);
        }
        else
        {
            scanf("%d %d",&a,&b);
            printf("%d\n",tree_query(a,b));
        }
    }
    return 0;
}