[bzoj] 2243 染色

原题

树链剖分板子题

线段树中直接记录染色段数，然后每个点记录左端点颜色和右端点颜色，以便于pushup

tre[i].sum=tre[i*2].sum+tre[i*2+1].sum-(tre[i*2].rc==tre[i*2+1].lc);

然后在爬树的过程中要单点查询颜色，用于判断是否-1。

线段树不要写跪，单点查询要记得pushdown#

#include<cstdio>
#include<algorithm>
#define N 100010
using namespace std;
int n,m,a[N],dfn[N],deep[N],head[N],tp[N],f[N],x,y,z,ref[N],son[N],size[N],cnt=1,t;
char j;
struct hhh
{
    int to,next;
}edge[2*N];
struct node
{
    int l,r,sum,lazy,lc,rc;
}tre[4*N];

int read()
{
    int ans=0,fu=1;
    char j=getchar();
    for (;(j<'0' || j>'9') && j!='-';j=getchar()) ;
    if (j=='-') fu=-1,j=getchar();
    for (;j>='0' && j<='9';j=getchar()) ans*=10,ans+=j-'0';
    return ans*fu;
}

void add(int u,int v)
{
    edge[cnt].to=v;
    edge[cnt].next=head[u];
    head[u]=cnt++;
}

void dfs1(int x,int fa,int dep)
{
    f[x]=fa;
    deep[x]=dep+1;
    size[x]++;
    int mx=0;
    for (int i=head[x],v;i;i=edge[i].next)
    {
	v=edge[i].to;
	if (v!=fa)
	{
	    dfs1(v,x,dep+1);
	    size[x]+=size[v];
	    if (size[v]>mx) mx=size[v],son[x]=v;
	}
    }
}

void dfs2(int x,int top)
{
    tp[x]=top;
    dfn[x]=++t;
    ref[t]=x;
    if (son[x]) dfs2(son[x],top);
    for (int i=head[x],v;i;i=edge[i].next)
    {
	v=edge[i].to;
	if (v!=son[x] && v!=f[x])
	    dfs2(v,v);
    }
}

void build(int i,int l,int r)
{
    tre[i].l=l;
    tre[i].r=r;
    tre[i].lazy=-1;
    if (l==r)
    {
	tre[i].lc=tre[i].rc=a[ref[l]];
	tre[i].sum=1;
	return ;
    }
    int mid=(l+r)>>1;
    build(i*2,l,mid);
    build(i*2+1,mid+1,r);
    tre[i].lc=tre[i*2].lc;
    tre[i].rc=tre[i*2+1].rc;
    tre[i].sum=tre[i*2].sum+tre[i*2+1].sum-(tre[i*2].rc==tre[i*2+1].lc);
}

void push(int x)
{
    if (tre[x].l==tre[x].r) return ;
    if (tre[x].lazy==-1) return ;
    tre[x*2].lazy=tre[x*2+1].lazy=tre[x].lazy;
    tre[x*2].sum=tre[x*2+1].sum=1;
    tre[x*2].lc=tre[x*2].rc=tre[x*2+1].lc=tre[x*2+1].rc=tre[x].lazy;
    tre[x].lazy=-1;
}

void modify(int i,int l,int r,int p)
{
    if (l>r) swap(l,r);
    if (tre[i].l==l && tre[i].r==r)
    {
	tre[i].lazy=p;
	tre[i].sum=1;
	tre[i].lc=tre[i].rc=p;
	return ;
    }
    push(i);
    int mid=(tre[i].l+tre[i].r)>>1;
    if (l>mid) modify(i*2+1,l,r,p);
    else if (r<=mid) modify(i*2,l,r,p);
    else
    {
	modify(i*2,l,mid,p);
	modify(i*2+1,mid+1,r,p);
    }
    tre[i].lc=tre[i*2].lc;
    tre[i].rc=tre[i*2+1].rc;
    tre[i].sum=tre[i*2].sum+tre[i*2+1].sum-(tre[i*2].rc==tre[i*2+1].lc);
}

int query(int i,int l,int r)
{
    if (l>r) swap(l,r);
    if (tre[i].l==l && tre[i].r==r) return tre[i].sum;
    push(i);
    int mid=(tre[i].l+tre[i].r)>>1;
    if (l>mid) return query(i*2+1,l,r);
    else if (r<=mid) return query(i*2,l,r);
    else return query(i*2,l,mid)+query(i*2+1,mid+1,r)-(tre[i*2].rc==tre[i*2+1].lc);
}

int ask(int i,int x)
{
    if (tre[i].l==x) return tre[i].lc;
    if (tre[i].r==x) return tre[i].rc;
    push(i);
    int mid=(tre[i].l+tre[i].r)>>1;
    if (x<=mid) return ask(i*2,x);
    else return ask(i*2+1,x);
}

void pathmodify(int u,int v,int p)
{
    while (tp[u]!=tp[v])
    {
	if (deep[tp[u]]<deep[tp[v]]) swap(u,v);
	modify(1,dfn[u],dfn[tp[u]],p);
	u=f[tp[u]];
    }
    modify(1,dfn[u],dfn[v],p);
}

int pathquery(int u,int v)
{
    int ans=0;
    while (tp[u]!=tp[v])
    {
	if (deep[tp[u]]<deep[tp[v]]) swap(u,v);
	ans+=query(1,dfn[u],dfn[tp[u]]);
	if (ask(1,dfn[tp[u]]) == ask(1,dfn[f[tp[u]]])) ans--;
	u=f[tp[u]];
    }
    ans+=query(1,dfn[u],dfn[v]);
    return ans;
}

int main()
{
    n=read();
    m=read();
    for (int i=1;i<=n;i++) a[i]=read();
    for (int i=1;i<n;i++)
    {
	x=read();
	y=read();
	add(x,y);
	add(y,x);
    }
    dfs1(1,0,0);
    dfs2(1,1);
    build(1,1,n);
    getchar();
    while(m--)
    {
	j=getchar();
	x=read();
	y=read();
	if (j=='C')
	{
	    z=read();
	    pathmodify(x,y,z);
	}
	else printf("%d\n",pathquery(x,y));
	getchar();
    }
    return 0;
}