[数据结构与算法]CF1551B2 Wonderful Coloring - 2 题解

本题我们首先用桶扫一遍序列，由于需要存出下标，所以我们把桶定义成vector形式。接下来扫一遍桶。由于方便，我们将每种颜料所涂的下标用vector来存

• 若桶中数字数量>=k，则将所有k个容器都加入一个，剩下来的就没办法了
• 若<k，则先把其暂存在另一个序列里，我们称之为“暂存序列”，长度为cnt

扫完我们遍历暂存序列，但是由于要满足每一种颜色染的字母的数量全部相等的条件，所以我们可以只遍历到cnt-cnt%k，感性理解就是最后一个不大于cnt且是k倍数的数。这样我们就依次放入容器就行了，我们就一个一个按顺序放，即为i%k+1。

暂存序列搞定之后，剩下的就是把所有容器存的数下标取出来，把结果数组的相应位置赋上答案即可。

代码实现：

#include <iostream>
#include <cstdio>
#include <vector>
#include <cstring>

using namespace std;

const int MAXN = 2e5 + 100;
int T, n, k, a[MAXN], b[MAXN], cnt, p[MAXN], vis[MAXN];
vector<int> bucket[MAXN], vec[MAXN];

int main() {
    scanf("%d", &T);
    while(T --) {
        scanf("%d%d", &n, &k);
        int cnt = 0;
        for(int i=1; i<=n; i++) {
            scanf("%d", &a[i]);
            bucket[a[i]].push_back(i);
        }
        for(int i=1; i<=n; i++) {
            if(bucket[a[i]].size() >= k && !vis[a[i]]) {
                for(int j=0; j<k; j++) {
                    vec[j+1].push_back(bucket[a[i]][j]);
                }
                vis[a[i]] = 1;
            } else if(bucket[a[i]].size() > 0 && !vis[a[i]]) {
                for(int j=0; j<bucket[a[i]].size(); j++) {
                    b[++cnt] = bucket[a[i]][j];
                }
                vis[a[i]] = 1;
            }
        }
        for(int i=1; i<=cnt-cnt%k; i++) {
            vec[i%k+1].push_back(b[i]); // 剩下的依次涂上各个颜色即可 无需考虑数字重复问题
        }
        for(int i=1; i<=k; i++) {
            for(int j=0; j<vec[i].size(); j++) {
                int id = vec[i][j];
                p[id] = i;
            }
        }
        for(int i=1; i<=n; i++) {
            printf("%d ", p[i]);
            // 小优化：memset比较慢 通过这种方式可以省一些常数
            bucket[i].clear();
            vec[i].clear();
            p[i] = 0;
            b[i] = 0;
            vis[i] = 0;
        }
        printf("\n");
    }
    return 0;
}