[백준] 20303 할로윈의 양아치 (C++)

링크: https://www.acmicpc.net/problem/20303

20303번: 할로윈의 양아치

문제

Trick or Treat!!

10월 31일 할로윈의 밤에는 거리의 여기저기서 아이들이 친구들과 모여 사탕을 받기 위해 돌아다닌다. 올해 할로윈에도 어김없이 많은 아이가 할로윈을 즐겼지만 단 한 사람, 일찍부터 잠에 빠진 스브러스는 할로윈 밤을 즐길 수가 없었다. 뒤늦게 일어나 사탕을 얻기 위해 혼자 돌아다녀 보지만 이미 사탕은 바닥나 하나도 얻을 수 없었다.

단단히 화가 난 스브러스는 거리를 돌아다니며 다른 아이들의 사탕을 빼앗기로 마음을 먹는다. 다른 아이들보다 몸집이 큰 스브러스에게 사탕을 빼앗는 건 어렵지 않다. 또한, 스브러스는 매우 공평한 사람이기 때문에 한 아이의 사탕을 뺏으면 그 아이 친구들의 사탕도 모조리 뺏어버린다. (친구의 친구는 친구다?!)

사탕을 빼앗긴 아이들은 거리에 주저앉아 울고 K명 이상의 아이들이 울기 시작하면 울음소리가 공명하여 온 집의 어른들이 거리로 나온다. 스브러스가 어른들에게 들키지 않고 최대로 뺏을 수 있는 사탕의 양을 구하여라.

스브러스는 혼자 모든 집을 돌아다녔기 때문에 다른 아이들이 받은 사탕의 양을 모두 알고 있다. 또한, 모든 아이는 스브러스를 피해 갈 수 없다.

입력

첫째 줄에 정수 N, M, K가 주어진다. N은 거리에 있는 아이들의 수, M은 아이들의 친구 관계 수, K는 울음소리가 공명하기 위한 최소 아이의 수이다. (1 ≤ N ≤ 30,000, 0 ≤ M ≤ 100,000, 1 ≤ k ≤ min{N ,3,000})

둘째 줄에는 아이들이 받은 사탕의 수를 나타내는 정수 c1,c2,⋯,cN이 주어진다. (1 ≤ ci ≤ 10,000)

셋째 줄부터 M개 줄에 갈쳐 각각의 줄에 정수 a, b가 주어진다. 이는 a와 b가 친구임을 의미한다. 같은 친구 관계가 두 번 주어지는 경우는 없다. (1≤ a, b ≤ N, a ≠ b)

출력

스브러스가 어른들에게 들키지 않고 아이들로부터 뺏을 수 있는 최대 사탕의 수를 출력한다.

---

풀이

스브러스는 한 아이의 사탕을 뺏을 때 그 아이의 친구들의 사탕도 모조리 뺏어버리기 때문에

BFS를 통해 아이를 선택하였을 때 우는 아이들의 수와 사탕의 수를 저장하였다.

그런데 무턱대로 30,000 × 30,000의 이차원 배열을 만들어 친구 정보를 저장하고 탐색을 하였더니 메모리 초과가 발생하였고, vector를 통해 BFS를 진행하였다.

 

예제 입력 1의 우는 아이들의 수와 사탕의 수 정보는 다음과 같다

(2, 13) → 1, 3번 아이
(4, 26) → 2, 5, 6, 10번 아이 
(2, 24) → 4, 9번 아이
(2, 33) → 7, 8번 아이

 

이렇게 정보를 저장하고 나면 여기에 배낭 문제를 적용하여 DP로 문제를 풀면 된다.

dp[i][j]에 우는 아이가 j명이고  i번째 그룹의 아이들을 적용했을 때 사탕의 최솟값을 저장하면 되고

예제 입력 1의 경우는 다음과 같다.

 

그런데 위와 같이 풀었을 때 사용한 메모리가 356544KB로 메모리 제한의 1/3정도였다.

아마 Union Find를 사용하여 우는 아이와 사탕의 수 정보를 저장한다면 더욱 효율적일 것이다.

코드

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

int candy[30001];
vector<int> friends[30001];
bool visited[30001];
int dp[30001][3001];

int n, m, k;

// 우는 아이들의 수와 사탕의 수를 BFS로 탐색해 반환해주는 함수
pair<int, int> bfs(int a){
    int numFriends = 0;
    int candies = 0;
    queue<int> q;
    q.push(a);
    visited[a] = true;

    while(!q.empty()){
        int child = q.front();
        q.pop();
        numFriends++;
        candies += candy[child];

        for (int i = 0; i < friends[child].size(); i++){
            int next = friends[child][i];
            if (!visited[next]){
                q.push(next);
                visited[next] = true;
            }
        }
    }
    return make_pair(numFriends, candies);
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(NULL); cout.tie(NULL);

    cin >> n >> m >> k;
    for (int i = 1; i <= n; i++){
        cin >> candy[i];
    }
    for (int i = 1; i <= m; i++){
        int a, b;
        cin >> a >> b;
        friends[a].push_back(b);
        friends[b].push_back(a);
    }

    // 우는 아이의 수와 사탕의 수 정보를 vec에 저장
    vector<pair<int, int>> vec;
    for (int i = 1; i <= n; i++){
        if (!visited[i]){
            vec.push_back(bfs(i));
        }
    }

    // dp로 최대 사탕의 수를 구하기 (배낭 문제)
    for (int i = 1; i <= vec.size(); i++){
        for (int j = 1; j < vec[i-1].first; j++){
            dp[i][j] = dp[i-1][j];
        }
        for (int j = vec[i-1].first; j < k; j++){
            dp[i][j] = max(dp[i-1][j], vec[i-1].second+dp[i-1][j-vec[i-1].first]);
        }
    }

    cout << dp[vec.size()][k-1];
}