线段树+Dfs序【CF620E】New Year Tree

Description

你有一棵以1为根的有根树，有n个点，每个节点初始有一个颜色c[i]。

有两种操作：

1 v c 将以v为根的子树中所有点颜色更改为c

2 v 查询以v为根的子树中的节点有多少种不同的颜色

Input

第一行,两个整数\(n,m\),分别代表有\(n\)个节点和\(m\)个操作。

第二行,共\(n\)个整数,代表每个节点的初始颜色\(c[i]\)

接下来\(n-1\)行,描述一条边。

接下来\(m\)行,代表每个操作。

Output

对于每个询问操作,输出一行。

刚开始以为是树剖?

结果发现只需要对每个子树操作。

线段树维护\(dfs\)序。

对于颜色呢？发现\(c[i] \leq 60\)

开$long \ long $可以压成一个数。

因此我们将颜色压缩即可。

记得开$long \ long $

虽然没出第二个样例,但我切了

代码

#include<cstdio>
#include<iostream>
#include<algorithm>
#define int long long 
#define R register

using namespace std;

const int gz=4e5+8;

inline void in(int &x)
{
	int f=1;x=0;char s=getchar();
	while(!isdigit(s)){if(s=='-')f=-1;s=getchar();}
	while(isdigit(s)){x=x*10+s-'0';s=getchar();}
	x*=f;
}

int head[gz],tot;

struct cod{int u,v;}edge[gz<<1];

inline void add(R int x,R int y)
{
	edge[++tot].u=head[x];
	edge[tot].v=y;
	head[x]=tot;
}

int dfn[gz],fdfn[gz],idx,size[gz];

void dfs(R int u,R int fa)
{
	dfn[u]=++idx;fdfn[idx]=u;size[u]=1;
	for(R int i=head[u];i;i=edge[i].u)
	{
		if(edge[i].v==fa)continue;
		dfs(edge[i].v,u);
		size[u]+=size[edge[i].v];
	}
}

int tr[gz<<2],c[gz],n,m;

bool tg[gz<<2];

#define ls o<<1
#define rs o<<1|1

inline void up(R int o)
{
	tr[o]=(tr[ls] | tr[rs]);
}

void build(R int o,R int l,R int r)
{
	if(l==r)
	{
		tr[o]=(1LL<<c[fdfn[l]]);
		return ;
	}
	R int mid=(l+r)>>1;
	build(ls,l,mid);
	build(rs,mid+1,r);
	up(o);
}

inline void down(R int o)
{
	if(tg[o])
	{
		tg[ls]=tg[rs]=tg[o];
		tr[ls]=tr[rs]=tr[o];
		tg[o]=false;
	}
}

void change(R int o,R int l,R int r,R int x,R int y,R int k)
{
	if(x<=l and y>=r){tr[o]=(1LL<<k);tg[o]=true;return;}
	down(o);
	R int mid=(l+r)>>1;
	if(x<=mid)change(ls,l,mid,x,y,k);
	if(y>mid)change(rs,mid+1,r,x,y,k);
	up(o);
}

int query(R int o,R int l,R int r,R int x,R int y)
{
	if(x<=l and y>=r)return tr[o];
	down(o);
	R int mid=(l+r)>>1;
	if(y<=mid)return query(ls,l,mid,x,y);
	else if(x>mid)return query(rs,mid+1,r,x,y);
	return (query(ls,l,mid,x,mid) | query(rs,mid+1,r,mid+1,y));
}

#define lowbit(o) o&-o 

inline int tquery(R int v)
{
	R int k=query(1,1,n,dfn[v],dfn[v]+size[v]-1);
	R int cnt=0;
	while(k) k-=lowbit(k),cnt++;
	return cnt;
}

signed main()
{
	in(n);in(m);
	for(R int i=1;i<=n;i++)in(c[i]);
	for(R int i=1,x,y;i<n;i++)
	{
		in(x),in(y);
		add(x,y);
		add(y,x);
	}
	dfs(1,0);build(1,1,n);
	for(R int i=1,opt,v,c;i<=m;i++)
	{
		in(opt);
		switch(opt)
		{
			case 1:in(v),in(c);change(1,1,n,dfn[v],dfn[v]+size[v]-1,c);break;
			case 2:in(v);printf("%lld\n",tquery(v));break;
		}
	}
}