题解：P10531 [XJTUPC2024] 圣诞树

喵仔牛奶 · 2024-05-20 21:51:58 · 题解

Solution

模仿 CF1946C 的做法，向下递归贪心，若当前连通块内颜色数 \geq k 就切断与父亲的边，否则与父亲的连通块合并。由于连通块没有权值，所以这样做是对的。

至于维护连通块颜色数，使用 set 启发式合并即可。复杂度 \mathcal{O}(n\log^2n)。

Code

#include <bits/stdc++.h>
#define REP(i, l, r) for (int i = (l); i <= (r); ++ i)
#define DEP(i, r, l) for (int i = (r); i >= (l); -- i)
#define fi first
#define se second
#define pb emplace_back
#define mems(x, v) memset((x), (v), sizeof(x))
using namespace std;
namespace Milkcat {
    typedef long long LL;
    typedef pair<LL, LL> pii;
    const int N = 1e6 + 5;
    int n, k, rs, u, v, c[N];
    vector<int> G[N]; set<int> s[N];
    void dfs(int u, int fa) {
        s[u].insert(c[u]);
        for (int v : G[u]) {
            if (v == fa) continue;
            dfs(v, u);
            if (s[u].size() < s[v].size()) s[u].swap(s[v]);
            for (int x : s[v]) s[u].insert(x);
        }
        if (s[u].size() >= k) s[u].clear(), rs ++;
    }
    int main() {
        cin >> n >> k;
        REP(i, 1, n) cin >> c[i];
        REP(i, 2, n) cin >> u >> v, G[u].pb(v), G[v].pb(u);
        dfs(1, 0);
        cout << rs << '\n';
        return 0;
    }
}
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    int T = 1;
    while (T --) Milkcat::main();
    return 0;
}