[백준 20303] 할로윈의 양아치

 

 

 알고리즘 몇 가지를 섞어서 접근해야 하는 재밌는 문제다. 딱 학부수준에서 다루는 알고리즘들을 적절하게 도입해서 해결할 수 있는데, 나 같은 경우에는 Union Find와 0-1 Knapsack Problem을 응용해서 해결했다. 

 

 문제를 천천히 읽어보면, 한 아이의 사탕을 뺏는 것은 친구관계로 연결된 모든 아이의 사탕을 한 번에 뺏는 것과 같다. 따라서 친구관계로 연결된 아이들의 집합을 찾아서 미리 소속된 아이 수와 아이들이 가진 사탕의 합을 찾아두면 빠르게 구할 수 있다. BFS같은 그래프 탐색을 이용해도 되고, 나는 Union Find를 사용했다.

 

 그렇게 친구끼리 묶은 집합 정보를 알아내고 나면, 이제 각 집합을 최적으로 선택해서 제한 아이 수 이내로 최대 사탕 수를 획득해야 함을 알 수 있다. 만약 브루트 포스로 모든 경우를 따지려고 한다면, 모든 아이가 친구관계가 아닐 경우에 2^30000 번 계산해야 하므로 시간 초과가 날 것이다.

 이 문제의 형태는 0-1 Knapsack Problem과 같다. K 값이 무게 상한, 아이 수가 각 아이템의 무게, 아이들이 가진 사탕의 합이 아이템의 가치라고 생각하면 된다. 따라서 Branch and Bound나 DP를 사용해서 접근해볼 수 있다. 

 나는 DP를 사용했고, 그 연산 수는 (집합의 수 * K)와 같으므로 최대 90,000,000회 가량 연산이 필요할 수 있다. 시간 제한이 1초인데 조금 빠듯한 느낌이 있기 때문에, 상대적으로 여유로운 공간 복잡도를 조금 희생하더라도 시간 복잡도에 좀 더 비중을 두고 작성하는게 좋다(그래서 분리 집합을 이용해 친구 집합 정보를 구축하는 과정에서 Red-Black 트리 기반 unordered_map을 이용했다가, 그냥 모조리 vector로 변경했다).

 풀면서 불필요한 로직이나, 부적절한 자료구조를 이용한다든지 반복문 내에서 메모리 locality를 신경쓰지 않거나 하면 시간 초과가 나겠다싶었다.

 

 DP 로직은 0-1 Knapsack Problem과 같아서 따로 쓸 건 없을 것 같다. 아이템(친구 집합)의 인덱스와 무게 제한(울음 공명 제한)을 가지고 DP 테이블을 만들어서, 아이템을 집어넣을 수 있는지 확인하고 넣을 수 있다면 넣는 경우(이전 인덱스 행에서 현재 탐색 중 무게 - 넣으려는 아이템의 무게 열)와 넣지 않는 경우(이전 인덱스 행의 현재 탐색 중 무게 열) 중 더 큰 값을 넣는 작업의 반복이다.

 DP 테이블을 완성한 이후에는, 무게 한도가 K-1이고 마지막 아이템까지 고려한 그 값을 출력해주면 된다.

 

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int root[30001];

int find(int x) {
    if(x == root[x]) {
        return x;
    } else {
        return root[x] = find(root[x]);
    }
}

void makeUnion(int x, int y) {
    int rootX = find(x);
    int rootY = find(y);

    if(rootX == rootY) return;
    if(rootX < rootY) {
        root[rootY] = rootX;
    } else {
        root[rootX] = rootY;
    }
}

int main() {
    cin.tie(nullptr);
    ios_base::sync_with_stdio(false);

    int n, m, k;
    cin >> n >> m >> k;

    vector<int> candyList(n+1);

    for(int i=1; i<=n; i++) {
        cin >> candyList[i];
        root[i] = i;
    }

    int child1, child2;
    for(int i=0; i<m; i++) {
        cin >> child1 >> child2;
        makeUnion(child1, child2);
    }

    vector<pair<int, int> > childUnionList; // { candySum, childNum }
    vector<int> candySum(n+1, 0);
    vector<int> childCount(n+1, 0);

    for(int i=1; i<=n; i++) {
        int unionNum = find(i);
        candySum[unionNum] += candyList[i];
        childCount[unionNum]++;
    }
    for(int i=1; i<=n; i++) {
        if(childCount[i]>0) {
            childUnionList.push_back({candySum[i], childCount[i]});
        }
    }

    vector<vector<int> > dp(childUnionList.size()+1, vector<int>(k, 0));

    for(int i=1; i<=childUnionList.size(); i++) {
        for(int j=1; j<k; j++) {
            if(childUnionList[i-1].second <= j) {
                dp[i][j] = max(childUnionList[i-1].first+dp[i-1][j-childUnionList[i-1].second], dp[i-1][j]);
            } else {
                dp[i][j] = dp[i-1][j];
            }
        }
    }
    
    cout << dp[childUnionList.size()][k-1];
}