五彩雪牛(Bessie's Snow Cow P) 题解

思路分析

按题目要求使用对应数据结构与算法模拟/递推求解，注意数据范围边界与时间复杂度。

参考代码

#include <iostream>
#include <vector>
#include <set>
using namespace std;

struct RangeBIT {
    int n;
    vector<long long> b1, b2;
    RangeBIT(int n) : n(n), b1(n + 2, 0), b2(n + 2, 0) {}
    void add(int idx, long long v) {
        long long v2 = v * idx;
        for (int i = idx; i <= n; i += i & -i) {
            b1[i] += v;
            b2[i] += v2;
        }
    }
    void rangeAdd(int l, int r, long long v) {
        add(l, v);
        add(r + 1, -v);
    }
    long long prefixSum(int idx) const {
        long long s1 = 0, s2 = 0;
        for (int i = idx; i > 0; i -= i & -i) {
            s1 += b1[i];
            s2 += b2[i];
        }
        return (idx + 1) * s1 - s2;
    }
    long long rangeSum(int l, int r) const {
        return prefixSum(r) - prefixSum(l - 1);
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, q;
    if (!(cin >> n >> q)) return 0;
    vector<vector<int>> adj(n + 1);
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    vector<int> dfn(n + 1, 0), sz(n + 1, 1), par(n + 1, 0), order;
    order.reserve(n);
    vector<int> stk = {1};
    while (!stk.empty()) {
        int u = stk.back();
        stk.pop_back();
        dfn[u] = (int)order.size() + 1;
        order.push_back(u);
        for (int i = (int)adj[u].size() - 1; i >= 0; --i) {
            int v = adj[u][i];
            if (v == par[u]) continue;
            par[v] = u;
            stk.push_back(v);
        }
    }
    for (int i = n - 1; i >= 0; --i) {
        int u = order[i];
        if (par[u]) sz[par[u]] += sz[u];
    }
    RangeBIT bit(n);
    vector<set<pair<int, int>>> colSet(100005);
    for (int i = 0; i < q; ++i) {
        int op;
        cin >> op;
        if (op == 1) {
            int x, c;
            cin >> x >> c;
            auto& st = colSet[c];
            auto it = st.upper_bound({dfn[x], n + 1});
            if (it != st.begin()) {
                auto prevIt = prev(it);
                int anc = prevIt->second;
                if (dfn[anc] + sz[anc] - 1 >= dfn[x]) continue;
            }
            it = st.lower_bound({dfn[x], 0});
            while (it != st.end() && it->first <= dfn[x] + sz[x] - 1) {
                int y = it->second;
                bit.rangeAdd(dfn[y], dfn[y] + sz[y] - 1, -1);
                it = st.erase(it);
            }
            st.insert({dfn[x], x});
            bit.rangeAdd(dfn[x], dfn[x] + sz[x] - 1, 1);
        } else {
            int x;
            cin >> x;
            cout << bit.rangeSum(dfn[x], dfn[x] + sz[x] - 1) << "\n";
        }
    }
    return 0;
}