P2486 [SDOI2011]染色树剖

た入场券 2024-04-18 12:44 65阅读 0赞

刚刚还给征宇讲：我他妈最讨厌这样贼长的代码了，思路很清晰，样例数据也能过，但是就是wa

幸福来得有点猝不及防~

一发AC

#include<bits/stdc++.h>
    using namespace std;
    const int M=1e5+50;
    struct Edge{int to,next;}e[M<<1];
    struct Node{int l,r,sum,lc,rc,v;}t[M<<2];
    int n,m,cnt,head[M<<1],a[M];
    int tot,dep[M],siz[M],son[M],fa[M],id[M],rk[M],top[M];
    
    void add(int x,int y);
    void dfs1(int x,int f);
    void dfs2(int x,int tp);
    void push_up(int k);
    void build(int k,int l,int r);
    void down(int k);
    void update(int k,int l,int r,int c);
    void updates();
    int ask(int k,int l,int r);
    int ask1(int k,int in);
    void query();
    
    int main(){
    	memset(head,-1,sizeof(head));cnt=1;
    	scanf("%d%d",&n,&m);
    	for(int i=1;i<=n;i++)scanf("%d",&a[i]);
    	for(int x,y,i=1;i<n;i++)scanf("%d%d",&x,&y),add(x,y);
    	dfs1(1,0);
    	dfs2(1,1);
    	build(1,1,n);
    	char op[3];
    	while(m--){
    		scanf("%s",op);
    		if(op[0]=='C')updates();
    		else query();
    	}
    	return 0;
    }
    
    
    void add(int x,int y){
    	e[cnt].to=y;
    	e[cnt].next=head[x];
    	head[x]=cnt++;
    	e[cnt].to=x;
    	e[cnt].next=head[y];
    	head[y]=cnt++;
    }
    void dfs1(int x,int f){
    	fa[x]=f;
    	dep[x]=dep[f]+1;
    	siz[x]=1;
    	for(int to,i=head[x];~i;i=e[i].next){
    		to=e[i].to;
    		if(to==f)continue;
    		dfs1(to,x);
    		siz[x]+=siz[to];
    		if(siz[ son[x] ]<siz[to]) son[x]=to;
    	}
    }
    void dfs2(int x,int tp){
    	top[x]=tp;
    	id[x]=++tot;
    	rk[tot]=a[x];
    	if(!son[x])return ;
    	dfs2(son[x],tp);
    	for(int i=head[x];~i;i=e[i].next)
    		if(!id[e[i].to])
    			dfs2(e[i].to,e[i].to);
    }
    void push_up(int k){
    	t[k].sum=t[k*2].sum+t[k*2+1].sum;
    	t[k].lc=t[k*2].lc;
    	t[k].rc=t[k*2+1].rc;
    	if(t[k*2].rc==t[k*2+1].lc)t[k].sum--;
    }
    void build(int k,int l,int r){
    	t[k].l=l;t[k].r=r;t[k].v=0;
    	if(l==r){
    		t[k].lc=t[k].rc=rk[l];t[k].sum=1;
    		return ;
    	}
    	int mid=(l+r)>>1;
    	build(k*2,l,mid);
    	build(k*2+1,mid+1,r);
    	push_up(k);
    }
    
    void down(int k){
    	if(t[k].v){
    		t[k*2].lc=t[k*2].rc=t[k*2].v=t[k].v;
    		t[k*2+1].lc=t[k*2+1].rc=t[k*2+1].v=t[k].v;
    		t[k*2].sum=t[k*2+1].sum=1;
    		t[k].v=0;
    	}
    }
    void update(int k,int l,int r,int c){
    	if(l<=t[k].l&&t[k].r<=r){
    		t[k].sum=1;
    		t[k].lc=t[k].rc=t[k].v=c;
    		return ;
    	}
    	int  mid=(t[k].l+t[k].r)>>1;
    	down(k);
    	if(l<=mid)update(k*2,l,r,c);
    	if(mid<r)update(k*2+1,l,r,c);
    	push_up(k);
    }
    void updates(){
    	int x,y,c;scanf("%d%d%d",&x,&y,&c);
    	while(top[x]!=top[y]){
    		if(dep[top[x]]<dep[top[y]])swap(x,y);
    		update(1,id[top[x]],id[x],c);
    		x=fa[top[x]];
    	}
    	if(id[x]>id[y])swap(x,y);
    	update(1,id[x],id[y],c);
    }
    int ask(int k,int l,int r){
    	if(l<=t[k].l&&t[k].r<=r)return t[k].sum;
    	down(k);
    	int res=0,mid=(t[k].l+t[k].r)>>1;
    	if(l<=mid)res+=ask(k*2,l,r);
    	if(mid<r)res+=ask(k*2+1,l,r);
    	if(l<=mid&&mid<r&&t[k*2].rc==t[k*2+1].lc)res--;
    	return res;
    }
    int ask1(int k,int in){
    	if(t[k].l==t[k].r)return t[k].lc;
    	int  mid=(t[k].l+t[k].r)>>1;
    	down(k);
    	if(in<=mid) return ask1(k*2,in);
    	return ask1(k*2+1,in);
    }
    
    void query(){
    	int x,y,res=0,c1,c2;
    	scanf("%d%d",&x,&y);
    	while(top[x]!=top[y]){
    		if(dep[top[x]]<dep[top[y]])swap(x,y);
    		res+=ask(1,id[top[x]],id[x]);
    		if(ask1(1,id[fa[top[x]]])==ask1(1,id[top[x]]))res--;
    		x=fa[top[x]];
    	}
    	if(id[x]>id[y])swap(x,y);
    	res+=ask(1,id[x],id[y]);
    	printf("%d\n",res);	
    }