B2 - Wonderful Coloring - 2

問題

https://codeforces.com/contest/1551/problem/B2 (新しいタブで開く)

問題文

数列が与えられる。数列を以下の条件を満たすように塗る。

1. 各要素は、どれか 1 つの色で塗られているか、もしくは塗られていない状態にする。
2. 同じ色で塗られている要素は、互いに異なる。
3. 各色で塗られている要素数は同じ。
4. 上記条件を満たすよう、塗られている要素数を最大化する。

$k$ 色で塗るとき、どう塗ればいい？

制約

• $1 \le t \le 10000$
• $1 \le n \le 2 \times 10^5$
• $1 \le k \le n$
• $n$ の総和は $2 \times 10^5$ を超えない

サンプル

6
10 3
3 1 1 1 1 10 3 10 10 2
4 4
1 1 1 1
1 1
1
13 1
3 1 4 1 5 9 2 6 5 3 5 8 9
13 2
3 1 4 1 5 9 2 6 5 3 5 8 9
13 3
3 1 4 1 5 9 2 6 5 3 5 8 9

1 1 0 2 3 2 2 1 3 3
4 2 1 3
1
0 0 1 1 0 1 1 1 0 1 1 1 0
2 1 2 2 1 1 1 1 2 1 0 2 2
1 1 3 2 1 3 3 1 2 2 3 2 0

考察

数字の出現回数を数える。

$k$ 回以上出てくる数字は、$k$ 色で塗る。
それ未満しか出てこない数字は、重複しないようにうまく塗ればいい。

そのために、まず塗るべき場所を列挙して、塗る色は $1, 2, \cdots, k, 1, \cdots$ とぐるぐるさせればいい。

コード

https://codeforces.com/contest/1551/submission/183322324 (新しいタブで開く)

void solve() {
  int n, k;
  cin >> n >> k;
  vector<int> a(n);
  rep(i, n) cin >> a[i];
  map<int, vector<int>> mp;
  rep(i, n) {
    if (mp[a[i]].size() < k) mp[a[i]].push_back(i);
  }
  int shouldPaint = 0;
  for (auto [_, v] : mp) shouldPaint += v.size();
  shouldPaint = (shouldPaint / k) * k;
  vector<int> ans(n, 0);
  int cur = 0;
  for (auto [_, v] : mp) {
    for (int i : v) {
      ans[i] = cur + 1;
      cur++;
      cur %= k;
      shouldPaint--;
      if (shouldPaint <= 0) break;
    }
    if (shouldPaint <= 0) break;
  }
  rep(i, n) cout << ans[i] << " ";
  cout << endl;
}