CF-620-E-DFS序+线段树

620-E 题目大意

给定一颗\(n\)个节点的树，根节点为\(1\)，每个节点都有一个颜色\(c_i\)。有\(m\)次操作，操作分两种：

• \(1,x,c\)：将以\(x\)为根的子树上的节点颜色都变为\(c\)。
• \(2,x\)：查询以\(x\)为根的子树上所有节点的颜色数量。

其中\(1{\le}c,c_i{\le}60\)。

---

Solution

首先预处理出树的\(dfs\)序，把子树上的操作转换为区间操作，注意到颜色数量不超过\(60\)，可以用一个\(64\)位的整数来表示一个子树中所有的颜色。

转到\(dfs\)序上的操作后，操作一为把一段区间推平为\(c\)；操作二为询问区间中颜色数量，这里可以对区间中所有颜色取并，用位运算的或即可得到。至此，所有的操作都可以用懒标记线段树解决。

时间复杂度\(O(nlogn)\)，难点在是否了解前置知识\(dfs\)序和线段树。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N=4e5+10;
struct node{
	int l,r,tag;
	ll cnt;
}tr[N<<2];

void pushdown(int u){
	if(tr[u].tag){
		tr[u<<1].cnt=tr[u<<1|1].cnt=1LL<<tr[u].tag;
		tr[u<<1].tag=tr[u<<1|1].tag=tr[u].tag;
		tr[u].tag=0;
	}
}

void build(int u,int l,int r,vector<int> &p){
	tr[u]={l,r,0,1LL<<p[l]};
	if(l==r) return;
	int m=(l+r)>>1;
	build(u<<1,l,m,p);
	build(u<<1|1,m+1,r,p);
	tr[u].cnt=tr[u<<1].cnt|tr[u<<1|1].cnt;
}

void modify(int u,int l,int r,int k){
	if(l<=tr[u].l&&tr[u].r<=r){
		tr[u].cnt=1LL<<k;
		tr[u].tag=k;
		return;
	}
	pushdown(u);
	int m=(tr[u].l+tr[u].r)>>1;
	if(l<=m) modify(u<<1,l,r,k);
	if(r>m) modify(u<<1|1,l,r,k);
	tr[u].cnt=tr[u<<1].cnt|tr[u<<1|1].cnt;
}

ll query(int u,int l,int r){
	if(l<=tr[u].l&&tr[u].r<=r){
		return tr[u].cnt;
	}
	pushdown(u);
	int m=(tr[u].l+tr[u].r)>>1;
	ll res=0;
	if(l<=m) res|=query(u<<1,l,r);
	if(r>m) res|=query(u<<1|1,l,r);
	return res;
}

void solve(){
	int n,m;
	cin>>n>>m;
	vector<int> col(n+1),p(n+1);
	for(int i=1;i<=n;i++) cin>>col[i];
	vector<vector<int>> e(n+1);
	for(int i=1;i<n;i++){
		int x,y;
		cin>>x>>y;
		e[x].push_back(y);
		e[y].push_back(x);
	}
	vector<int> L(n+1),R(n+1);
	int dfn=0;
	function<void(int,int)> dfs=[&](int x,int fa){
		L[x]=++dfn,p[dfn]=col[x];
		for(auto y:e[x]){
			if(y==fa) continue;
			dfs(y,x);
		}
		R[x]=dfn;
	};
	dfs(1,0);
	build(1,1,n,p);
	while(m--){	
		int op,x,c;
		cin>>op>>x;
		if(op&1){
			cin>>c;
			modify(1,L[x],R[x],c);
		}else{
			ll res=query(1,L[x],R[x]);
			cout<<__builtin_popcountll(res)<<'\n';
		}
	}
}

int main(){
	ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
	int T=1;
	//cin>>T;
	while(T--){
		solve();
	}
	return 0;
}