BOJ 20303 : 할로윈의 양아치 [C++]

 

20303번: 할로윈의 양아치

 

생각보다 금방금방 아이디어가 떠올랐던 문제

 

친구 그룹 중 한 명이 뺏기면 친구 관계인 다른 친구들도 사탕을 모두 뺏긴다는 점에서 union find를 생각했고,

그룹 중 한 명만 뺏을 수는 없기 때문에 그룹의 사이즈를 알아야 되겠다 생각해 weighted union find를 선택했다.

 

근데 weighted union find의 코드가 정확하게 기억이 안 나서 예전에 봤던걸 다시 봤다.

 

처음에는 친구 그룹끼리 묶어주고, {사탕 개수, 그룹 크기}를 PQ에 넣어 그리디 하게 가져가면 되겠지 했다.

예제도 통과하길래 제출했는데 실패

 

'그럼 DP로 해야 하나?'라고 생각했고, K명을 넘지 않도록 사탕을 최대한 뺏아야 한다는 점에서 냅색을 생각했다.

 

    for (int i = 0; i< kid_group.size(); i++) {
        for (int j = 0; j < K; j++) {
            if (j >= kid_group[i].second) {
                dp[j] = std::max(dp[j], dp[j - kid_group[i].second] + kid_group[i].first);
            }
        }
    }

 

처음엔 이렇게 했는데 빼앗긴 아이들의 수를 오름차순으로 체크하니 지금 뺏는 그룹의 사탕 수가 중복으로 카운트되어 나와야 할 값보다 큰 값이 나왔다.

 

    for (int i = 0; i< kid_group.size(); i++) {
        for (int j = K - 1; j > 0; j--) {
            if (j >= kid_group[i].second) {
                dp[j] = std::max(dp[j], dp[j - kid_group[i].second] + kid_group[i].first);
            }
        }
    }

 

예전 어떤 문제에서 내림차순으로 구해 중복을 처리했던 것이 기억나 내림차순으로 바꿔주니 제대로 답이 나왔다. 

 

정리

1. weighted union find로 친구들끼리 그룹을 지어준다

2. 그룹끼리 합치며 사탕 개수도 합쳐준다. (그룹 크기는 weighted union find로 알아서 관리됨)

3. {사탕 개수, 그룹 크기}로 구성된 pair를 만들어 vector에 저장

4. k명의 사탕을 뺏았을 때 가질 수 있는 사탕의 최대치를 표현하는 1차원 DP 배열을 이용해 냅색을 조진다

 

 

코드

#include <iostream>
#include <vector>

#define PII std::pair<int, int>

int N, M, K;
int candies[30001];
int parent[30001];
int dp[3001];

int find_parent(int x) {
    if (parent[x] < 0) {
        return x;
    } else {
        return parent[x] = find_parent(parent[x]);
    }
}

int main() {
    std::cin >> N >> M >> K;

    for (int i = 1; i < N + 1; i++) {
        std::cin >> candies[i];
        parent[i] = -1;
    }

    int A, B;
    for (int i = 0; i < M; i++) {
        std::cin >> A >> B;

        int pA = find_parent(A);
        int pB = find_parent(B);

        if (pA == pB) {
            continue;
        }

        if (parent[pA] < parent[pB]) {
            candies[pA] += candies[pB];
            candies[pB] = 0;
            parent[pA] += parent[pB];
            parent[pB] = pA;
        } else {
            candies[pB] += candies[pA];
            candies[pA] = 0;
            parent[pB] += parent[pA];
            parent[pA] = pB;
        }
    }

    std::vector<PII> kid_group;
    for (int i = 1; i < N + 1; i++) {
        if (parent[i] < 0) {
            kid_group.push_back({candies[i], -parent[i]});
        }
    }

    for (int i = 0; i< kid_group.size(); i++) {
        for (int j = K - 1; j > 0; j--) {
            if (j >= kid_group[i].second) {
                dp[j] = std::max(dp[j], dp[j - kid_group[i].second] + kid_group[i].first);
            }
        }
    }

    std::cout << dp[K - 1] << "\n";

    return 0;
}