CF620E New Year Tree (线段树维护 dfs 序)

CF620E New Year Tree

题意：给出一棵 n 个节点的树，根节点为 1。每个节点上有一种颜色 ci​。m 次操作。操作有两种：

• 1 u c：将以 u 为根的子树上的所有节点的颜色改为 c。
• 2 u：询问以 u 为根的子树上的所有节点的颜色数量。
  
  1 <= c <= 60。

由于 c 的范围，可以用一个整数来表示每棵子树的状态。

将树上问题转化为序列问题，当然可以树剖，但本题只对子树操作，维护 in 为进树时的时间戳，out 为出树时的时间戳，再在 dfs 序上用线段树维护即可。

和 CF1916E 很像。

#include<bits/stdc++.h>
#define rep(i, a, b) for(int i = (a); i <= (b); ++ i)
#define per(i, a, b) for(int i = (a); i >= (b); -- i)
#define pb emplace_back
#define All(X) X.begin(), X.end()
#define ls x << 1
#define rs ls | 1
using namespace std;
using ll = long long;
constexpr int N = 4e5 + 5;

int n, m, c[N], in[N], out[N], timestamp;
vector<int> G[N];

struct Node {
	ll v, tag;
} t[N << 2];

void pushup(int x) {
	t[x].v = t[ls].v | t[rs].v;
}

void pushdown(int x) {
	if(t[x].tag) {
		t[ls].v = t[ls].tag = t[x].tag;
		t[rs].v = t[rs].tag = t[x].tag;
		t[x].tag = 0;
	}
}

void modify(int L, int R, int v, int x = 1, int l = 1, int r = n) {
	if(L <= l && r <= R) {
		t[x].v = t[x].tag = 1ll << v;
		return;
	}
	pushdown(x);
	int mid = l + r >> 1;
	if(mid >= L) modify(L, R, v, ls, l, mid);
	if(mid < R) modify(L, R, v, rs, mid + 1, r);
	pushup(x);
}

ll query(int L, int R, int x = 1, int l = 1, int r = n) {
	if(L <= l && r <= R) return t[x].v;
	pushdown(x);
	int mid = l + r >> 1;
	ll ret = 0;
	if(mid >= L) ret |= query(L, R, ls, l, mid);
	if(mid < R) ret |= query(L, R, rs, mid + 1, r);
	return ret;
}

void dfs(int x, int fa) {
	in[x] = ++ timestamp;
	modify(in[x], in[x], c[x]);
	for(int y : G[x]) {
		if(y != fa) {
			dfs(y, x);
		}
	}
	out[x] = timestamp;
}

int main() {
	ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);
	cin >> n >> m;
	rep(i, 1, n) cin >> c[i];
	rep(i, 2, n) {
		int x, y; cin >> x >> y;
		G[x].pb(y);
		G[y].pb(x);
	}
	dfs(1, 0);
	rep(i, 1, m) {
		int op, x; cin >> op >> x;
		if(op == 1) {
			int col; cin >> col;
			modify(in[x], out[x], col);
		}
		else cout << __popcount(query(in[x], out[x])) << '\n';
	}
	return 0;
}