Codeforces Round #734 (Div. 3)-B2. Wonderful Coloring - 2

문제 설명

폴은 정수 $a_1, a_2, ..., a_n$ 으로 이루어진 수열을 생각했다. 폴은 이 정수들을 $k$개 색깔의 분필로 칠하고 싶다. 다음의 조건을 만족하면서 색칠을 하면 환상적인 수열이라고 말한다.

1. 수열의 각 원소는 $k$개의 색깔중 하나만 색칠하거나 색칠하지 않아야 한다.
2. 같은 색깔로 색칠한 모든 두 원소는 서로 다른 값을 가져야한다.
3. 각 $k$번째 색깔로 칠한 원소의 개수는 모두 같아야 한다.
4. 위에서 언급한 3개의 조건을 만족하면서 색칠한 원소의 개수가 최대가 되어야 한다.

예를 들어 수열 $a = [3, 1, 1, 1, 10, 3, 10, 10, 2]$ 이고 $k = 3$ 인 경우 수열의 환상적인 색칠의 한가지 경우는 다음과 같다.

폴이 주어진 수열 $a$를 환상적인 색칠을 할 수 있게 도와주자.

Input
첫번째 줄에는 테스트 케이스의 수 $t (1 \le t \le 10000)$ 이 주어진다.

각 테스트 케이스의 첫번째 줄에는 주어진 수열의 길이와 색깔이 개수 $n$과 $k (1 \le n \le 2 \cdot 10^5, 1 \le k \le n)$ 이 주어진다.

각 테스트 케이스의 두번째 줄에는 정수 $n$ 개 $a_1, a_2, ..., a_n (1 \le a_i \le n)$ 이 주어진다.

Output
각 테스트케이스마다 환상적인 색칠을 한 수열의 모습을 출력해야 한다.

각 환상적인 색칠을 한 수열은 다음과 같은 조건을 만족하는 $n$개의 정수 $c_1, c_2, ..., c_n (0 \le c_i \le k)$를 출력해야 한다.

• $c_i = 0$ $i$번째 원소가 색칠이 안되어 있을때
• $c_i > 0$ $i$번째 원소가 $c_i$번째 색깔로 색칠되어 있을때

우리는 환상적인 색칠을 한 수열에서 각 색깔로 색칠된 원소의 개수를 최대화 하는 것이 목적이라는 것을 잊지 말아야한다.

Example
input
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

output
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

문제 접근

사용한 알고리즘: 그리디 알고리즘, 자료 구조, 구성적
걸린 시간 : 00:07
이것도 조금만 생각하면 풀 수 있다고 생각하면 큰 오산이다. 나는 생각보다 어려웠다. 그냥 B1문제의 일반적인 케이스인데, 색깔을 칠해야 하는 부분에서 구현이 꼬였다. 또한 맞왜틀이 너무 많이 나왔었는데, 대회가 끝나고 나서 editorial을 보고 내가 놓친 부분이 있었던 것 같다. 이 부분은 문제를 해설하면서 찬찬히 생각해보자.

기본적인 골격은 B1에서 해설한 부분과 같다. $k$보다 큰 경우는 정확히 $k$개 원소만 색칠을 해야하고 나머지 원소는 색칠을 하면 안된다. 그리고 $k$보다 작은 경우 그 원소의 총 합을 계산해서 $k$로 나눠서 숫자를 번갈아 가며 색칠을 하면 된다.

그런데 구현을 잘 못했다. 고수들은 구현을 어떻게 했을까? 각 숫자가 나타난 index를 2차원 벡터에 넣음으로 빈도수와 그 index의 위치를 모두 표현했다. 이쪽 부분이 자료구조와 관련된 부분인 것 같다.

그리고 위에서 언급한 방식대로 $k$보다 큰 경우는 $k$개만 ans벡터에 넣고, $k$보다 작은 경우는 그 개수 만큼 ans벡터에 넣는다. 이렇게 되면 같은 원소들은 무조건 연속적으로 ans벡터에 들어가기 때문에 $k$개의 색깔을 $1$, $2$, ... $k$으로 순차적으로 부여한다면 절대 원소끼리 겹칠일이 없다.

여기서 주의해야할 부분은 ans 안에 있는 원소의 개수가 $k$의 배수가 안되면 남은 숫자들은 색칠을 하면 안된다는 것이다. 왜냐하면 각 색깔로 색칠한 수의 개수가 모두 동일해야 하는데, 그냥 무지성으로 ans안에 있는 원소를 모두 순차적으로 색깔을 부여한다면 각 색깔로 색칠한 수의 개수가 동일 할 수가 없는 경우가 생길 수 있기 때문이다.

정답 코드

#include <iostream>
#include <vector>
using namespace std;

int main() {
    ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    int t; cin >> t;
    while(t--){
        int n, k; cin >> n >> k;
        vector<int> cnt[n + 1];
        for(int i = 0; i < n; i++){
            int num; cin >> num;
            cnt[num].push_back(i);
        }

        vector<int> color;
        for(int i = 1; i <= n; i++){
            if(cnt[i].size() >= k){
                for(int j = 0; j < k; j++) color.push_back(cnt[i][j]);
            }else{
                for(int j = 0; j < cnt[i].size(); j++) color.push_back(cnt[i][j]);
            }
        }

        vector<int>ans(n);
        for(int i = 0; i < k * (color.size() / k); i++) ans[color[i]] = i % k + 1;
        for(int x : ans) cout << x << " ";
        cout << "\n";
    }
    return 0;
}