[CF620E] New Year Tree

[CF620E] New Year Tree - 线段树,DFS序,位运算

Description

给出一棵 (n) 个节点的树，根节点为 (1)。每个节点上有一种颜色 (c_i)。(m) 次操作。操作有两种：

1. 1 u c：将以 (u) 为根的子树上的所有节点的颜色改为 (c)。
2. 2 u：询问以 (u) 为根的子树上的所有节点的颜色数量。
   (1le n,mle 4 imes 10^5)，(1le c_i,cle 60)

Solution

dfs 序转化为区间问题，线段树处理，每个结点用一个 int64 存颜色，存标记，需要有标记下传和结果上传，支持区间修改和区间询问

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

#define int long long

struct Tree
{
    vector<vector<int>> g;
    int n;
    Tree(int n) : n(n)
    {
        g.resize(n + 2);
    }

    void make(int p, int q)
    {
        g[p].push_back(q);
        g[q].push_back(p);
    }

    vector<int> dfn, fin;

    void solve()
    {
        int ind = 0;
        dfn.resize(n + 2);
        fin.resize(n + 2);
        function<void(int, int)> dfs = [&](int p, int from) -> void {
            dfn[p] = ++ind;
            for (auto q : g[p])
                if (q != from)
                {
                    dfs(q, p);
                }
            fin[p] = ind;
        };
        dfs(1, 0);
    }
};

struct SegmentTree
{
    int n;

    struct Node
    {
        int val;
        int tag;
        Node &operator<<(const Node &rhs)
        {
            if (rhs.tag)
            {
                val = rhs.tag;
                tag = rhs.tag;
            }
        }
        friend Node operator+(const Node &lhs, const Node &rhs)
        {
            return {lhs.val | rhs.val, 0};
        }
    };

    vector<Node> nodes;

    SegmentTree(int n) : n(n)
    {
        nodes.resize(4 * n + 8);
    }

    void pushup(int p)
    {
        nodes[p] = nodes[p * 2] + nodes[p * 2 + 1];
    }

    void pushdown(int p)
    {
        nodes[p * 2] << nodes[p];
        nodes[p * 2 + 1] << nodes[p];
        nodes[p].tag = 0;
    }

    void build(int p, int l, int r, const vector<int> &src)
    {
        if (l == r)
        {
            nodes[p] = {1ll << src[l], 0};
        }
        else
        {
            build(p * 2, l, (l + r) / 2, src);
            build(p * 2 + 1, (l + r) / 2 + 1, r, src);
            pushup(p);
        }
    }

    void Build(const vector<int> &src)
    {
        build(1, 1, n, src);
    }

    void modify(int p, int l, int r, int ql, int qr, int val)
    {
        if (l > qr || r < ql)
            return;
        if (l >= ql && r <= qr)
            nodes[p] << (Node){1ll << val, 1ll << val};
        else
        {
            pushdown(p);
            modify(p * 2, l, (l + r) / 2, ql, qr, val);
            modify(p * 2 + 1, (l + r) / 2 + 1, r, ql, qr, val);
            pushup(p);
        }
    }

    void Modify(int ql, int qr, int val)
    {
        modify(1, 1, n, ql, qr, val);
    }

    Node query(int p, int l, int r, int ql, int qr)
    {
        if (l > qr || r < ql)
            return {0, 0};
        if (l >= ql && r <= qr)
            return nodes[p];
        pushdown(p);
        return query(p * 2, l, (l + r) / 2, ql, qr) + query(p * 2 + 1, (l + r) / 2 + 1, r, ql, qr);
    }

    int Query(int ql, int qr)
    {
        return __builtin_popcountll(query(1, 1, n, ql, qr).val);
    }
};

signed main()
{
    ios::sync_with_stdio(false);

    int n, m;
    cin >> n >> m;

    vector<int> c(n + 2);
    for (int i = 1; i <= n; i++)
        cin >> c[i];

    SegmentTree seg(n);
    Tree tree(n);

    for (int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        tree.make(u, v);
    }

    tree.solve();

    auto &dfn = tree.dfn;
    auto &fin = tree.fin;

    vector<int> src(n + 2);
    for (int i = 1; i <= n; i++)
        src[dfn[i]] = c[i];

    seg.Build(src);

    for (int i = 1; i <= m; i++)
    {
        int t1, t2, t3;
        cin >> t1;
        if (t1 == 1)
        {
            cin >> t2 >> t3;
            int start = dfn[t2], end = fin[t2];
            seg.Modify(start, end, t3);
        }
        else
        {
            cin >> t2;
            int start = dfn[t2], end = fin[t2];
            cout << seg.Query(start, end) << endl;
        }
    }
}