【每日一题】New Year Tree 题解(树链剖分 线段树 状压)

Description

给一棵树，根为1，每个节点有一种颜色，有两种操作：

• 将该节点x及其子树的颜色都变为y
• 查询节点x及其子树有多少种颜色

Solution

颜色的范围为 (), 不妨用二进制上的每一位去代表每一种颜色是否存在
那么检查有多少种颜色只需要统计二进制的1个数，用函数 即可
由于是树上的操作，先进行树链剖分，将树上操作改为 序列上的区间操作
那么就转化为经典的线段树问题，用线段树维护并进行区间位运算即可。

Code

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
const int N = 4e5 + 5;
const int mod = 998244353;
vector<int> G[N];
int L[N], R[N], rk[N], a[N], tot;
struct node {
  int l, r;
  ll val;
  bool lazy;
}t[N << 2];
void dfs(int u, int par) {
  L[u] = ++tot;
  rk[tot] = u;
  for(auto &v : G[u]) {
    if(par == v) continue;
    dfs(v, u);
  }
  R[u] = tot;
}
void push_down(int rt) {
  if(t[rt].lazy) {
    t[rt << 1].lazy = t[rt << 1 | 1].lazy = t[rt].lazy;
    t[rt << 1].val = t[rt << 1 | 1].val = t[rt].val;
    t[rt].lazy = false;
  }
}
void push_up(int rt) {
  t[rt].val = t[rt << 1].val | t[rt << 1 | 1].val;
}
void build(int rt, int l, int r) {
  t[rt].l = l, t[rt].r = r;
  if(l == r) {
    t[rt].val |= (1LL << a[rk[l]]);
    return ;
  }
  int mid = t[rt].l + t[rt].r >> 1;
  build(rt << 1, l, mid);
  build(rt << 1 | 1, mid + 1, r);
  push_up(rt);
}
void update(int rt, int ql, int qr, int c) {
  if(ql <= t[rt].l && qr >= t[rt].r) {
    t[rt].val = (1LL << c);
    t[rt].lazy = true;
    return ;
  }
  push_down(rt);
  int mid = t[rt].l + t[rt].r >> 1;
  if(qr <= mid) {
    update(rt << 1, ql, qr, c);
  } else if(ql > mid) {
    update(rt << 1 | 1, ql, qr, c);
  } else {
    update(rt << 1, ql, mid, c);
    update(rt << 1 | 1, mid + 1, qr, c);
  }
  push_up(rt);
}
ll query(int rt, int ql, int qr) {
  if(ql <= t[rt].l && qr >= t[rt].r) {
    return t[rt].val;
  }
  push_down(rt);
  int mid = t[rt].l + t[rt].r >> 1;
  if(qr <= mid) {
    return query(rt << 1, ql, qr);
  } else if(ql > mid) {
    return query(rt << 1 | 1, ql, qr);
  } else {
    return (query(rt << 1, ql, mid) | query(rt << 1 | 1, mid + 1, qr));
  }
}
int main() {
  ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);
  int n, m; cin >> n >> m;
  for(int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  for(int i = 1; i <= n - 1; i++) {
    int u, v; cin >> u >> v;
    G[u].emplace_back(v);
    G[v].emplace_back(u);
  }
  dfs(1, 0);
  build(1, 1, n);
  for(int i = 1; i <= m; i++) {
    int op, x, y; cin >> op;
    if(op == 1) {
      cin >> x >> y;
      update(1, L[x], R[x], y);
    } else {
      cin >> x;
      cout << __builtin_popcountll(query(1, L[x], R[x])) << '\n';
    }
  }
  return 0;
}