[Algorithm]Baekjoon 20303번 :: 할로윈의 양아치

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

20303번: 할로윈의 양아치

#include <iostream>
#include <string>
#include <map>
#include <vector>
#include <algorithm>
#include <math.h>
#include <queue>
#include <functional>
#include <sstream>
#include <memory.h>
#include <deque>
#include <set>
#include <unordered_map>
#include <stack>
#include <numeric>
 
using namespace std;
 
 
int n, m, k;
int candyCounts[30001] = { 0 };
int groupCounts[30001] = { 0 };
int parents[30001] = { 0 };
int dp[30001][3001] = { 0 };
vector<pair<int, int>> groupes = { {0,0} };
 
int getParent(int node)
{
    if (node == parents[node]) return node;
    return parents[node] = getParent(parents[node]);
}
 
void unionParents(int nodeA, int nodeB)
{
    nodeA = getParent(nodeA);
    nodeB = getParent(nodeB);
 
    if (nodeA < nodeB)
    {
        parents[nodeB] = nodeA;
    }
    else
    {
        parents[nodeA] = nodeB;
    }
}
 
bool isSameParents(int nodeA, int nodeB)
{
    nodeA = getParent(nodeA);
    nodeB = getParent(nodeB);
 
    if (nodeA == nodeB) return true;
    return false;
}
 
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
 
    cin >> n >> m >> k;
 
    for (int i = 1; i <= n; ++i)
    {
        cin >> candyCounts[i];
        parents[i] = i;
        groupCounts[i] = 1;
    }
 
    for (int i = 0; i < m; ++i)
    {
        int a, b;
        cin >> a >> b;
 
        if (isSameParents(a, b)) continue;
 
        unionParents(a, b);
    }
 
    for (int i = 1; i <= n; ++i)
    {
        if (i != parents[i]) // 자기자신 아니면
        {
            int p = getParent(i);
            candyCounts[p] += candyCounts[i];
            groupCounts[p] += groupCounts[i];
        }
    }
 
    for (int i = 1; i <= n; ++i)
    {
        if (i != parents[i]) continue;
        
        groupes.push_back({ groupCounts[i], candyCounts[i] });
    }
 
    for (int i = 1; i < groupes.size(); ++i)
    {
        int groupCount = groupes[i].first;
        int groupCandy = groupes[i].second;
 
        for (int j = 1; j < k; ++j)
        {
            if (groupCount <= j) // 담을 수 있으면
            {
                dp[i][j] = max(dp[i - 1][j], dp[i - 1][j - groupCount] + groupCandy);
            }
            else
            {
                dp[i][j] = dp[i - 1][j];
            }
        }
    }
 
    cout << dp[groupes.size()-1][k-1];
}

 

Union-Find를 이용해서 각각의 서로소집합을 구한 후, k개 미만으로 최대의 캔디개수를 구하는 문제이다.

즉 각각의 집합들(원소개수, 원소들이 들고있는 캔디들의 합)으로, 배낭문제를 푼다고 생각하면 된다.

 

사실 푸는방법은 금방 알아냈고, 배낭풀이도 이제 바로바로 생각나서 코드작성이 가능한데 86번째줄 하나때문에 시간을 좀 먹었다. 해당 시점에서 여전히 parents[i]에 i의 최상위부모가 들어있다는게 보장되어지지 않기 때문에, getParent를 한번 더 호출해서 parents테이블을 업데이트해줘야 한다.