CF1551B2 Wonderful Coloring - 2

https://codeforces.com/problemset/problem/1579/D

题意

n个数字，分成k组，每组内不能有相同元素，每组的元素数量相同

求每组最多能放多少元素 并且记录每个数字放在哪一组了 没分组的元素输出0

Code

#include<bits/stdc++.h>
using namespace std;
#define int long long 
#define pii pair<int,int>
#define ar2 array<int,2>
#define ar3 array<int,3>
#define ar4 array<int,4>
#define endl '\n'
void cmax(int &a,int b){a=max(a,b);}
void cmin(int &a,int b){a=min(a,b);}
const int N=2e5+10,MOD=1e9+7,INF=0x3f3f3f3f,LINF=LLONG_MAX;
int a[N],n,k,ans[N];
unordered_map<int,int>freq,cnt;

struct node{
    int val;
    int id;
    
};

bool cmp(node x,node y){
    return x.val<y.val;
}

void solve() {
    cin>>n>>k;
    freq.clear();cnt.clear();
    memset(ans,0,sizeof ans);
    for(int i=1;i<=n;i++){
        cin>>a[i];
        freq[a[i]]++;
    }
    vector<node>v;//存储freq<k的元素
    for(int i=1;i<=n;i++){
        //当前赋值给ans的是cnt+1 如果cnt+1>k就赋值为0
        if(freq[a[i]]>=k){
            if(cnt[a[i]]+1>k) continue;
            ans[i]=++cnt[a[i]];
        }else if(freq[a[i]]<k){
            v.emplace_back((node){a[i],i});
        }
    }
    /*
    处理freq小于k的元素
    
    */
   sort(v.begin(),v.end(),cmp);
    int kk=0,m=v.size();
    for(int i=0;i<m;i++){
        if(m-i<k&&kk==0) break;
        ans[v[i].id]=++kk;
        if(kk+1>k) kk=0;
    }
    for(int i=1;i<=n;i++){
        cout<<ans[i]<<' ';
    }
    cout<<endl;
}   

signed main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);

    int t=1;
    cin>>t;
    while (t--) solve();

}

freq[a[i]]存储每个元素的出现次数

cnt[a[i]]存的是当前用了多少个a[i]        这样搞的目的是正好可以用++cnt[a[i]]来染色 标记组号

核心逻辑

先处理freq>=k的元素类别：

这类元素可以在所有k组里放一个 另外多余的就是0

同时把freq<k的元素种类单独找出来 用结构体数组v存起来 同时保留该元素在原数组中的下标

如何处理v里的元素：

先排序 这样所有相同的元素会在一起 

然后遍历v 依次给每个元素赋值 利用++kk 回滚地染色

何时跳出循环：

如果剩下的元素少于k个了 说明剩下的元素不能平均地分配给k组了 这样就会导致k个组中的元素个数不同 所以这时候就要break