洛谷 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全涂为一种颜色，颜色段数量为1t[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);elsereturn 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);elseupdata(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]);elseres+=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;
}