题解:P10531 [XJTUPC 2024] 圣诞树

· · 题解

思路:

考虑从叶子节点向上贪心,能选出一个圣诞树就选,能选就选肯定是不劣的。对每一个点都开一个 set 维护子树里有多少种颜色,如果颜色数大于等于 k 就将 set 清空。合并 set 时用启发式合并可以减小复杂度。

时间复杂度:O(n\log^2 n)。

代码:

#include<bits/stdc++.h>

using namespace std;
const int N=200010;
int n,k,c[N],ans;
vector<int> vec[N];
multiset<int> st2[N];
set<int> st[N];
void dfs(int u,int fa)
{
    st[u].insert(c[u]);
    for(auto i:vec[u]){
        if(i==fa) continue;
        dfs(i,u);
        if(st[i].size()>st[u].size()){
            swap(st[i],st[u]);
        }
        for(auto i:st[i]){
            st[u].insert(i);
        }
        st[i].clear();
    }
    if(st[u].size()>=k)
    {
        ans++;
        st[u].clear();
    }
}
int main(){
    cin>>n>>k;
    for(int i=1;i<=n;i++) cin>>c[i];
    for(int i=1;i<n;i++){
        int x,y;
        cin>>x>>y;
        vec[x].push_back(y);
        vec[y].push_back(x);
    }
    dfs(1,0);
    cout<<ans;
    return 0;
}