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

이 문제는 유니온 파인드, 배낭 문제 알고리즘을 이용해서 풀었다.

 

한 아이의 사탕을 뺏으면 친구들의 사탕도 모조리 뺏는다는 조건을 보고 유니온 파인드를 떠올렸다.

친구 그룹을 형성하여 총 사탕 수와 총 인원 수를 각각 candy, groupCount에 저장했다.

 

또한 각 그룹을 선택할지 말지에 대한 배낭 문제와 동일하다는 생각이 들어 배낭 문제의 알고리즘을 그대로 사용했다.

이 때 주의할 점은 dp[i - 1]에 대한 비교가 계속되므로 i = 0일 때 i = -1인 지점을 찾게 되어 오류가 발생할 수 있다.

따라서 i = 1부터 시작하도록 만들어 문제를 해결했다.

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

#define MAX 30000

vector<pair<int, int>> friends;
int groupCount[MAX + 1];
int N, M, K;
int candy[MAX + 1];
int parent[MAX + 1];
int dp[MAX + 1][3000];

void init() {
	for (int i = 1; i <= N; i++) {
		parent[i] = i;
	}
}

int find(int x) {
	if (parent[x] != x) {
		return parent[x] = find(parent[x]);
	}
	return x;
}

void merge(int a, int b) {
	a = find(a);
	b = find(b);
	if (a == b) {
		return;
	}
	if (a < b) {
		parent[b] = a;
		candy[a] += candy[b];
		candy[b] = 0;
	} else {
		parent[a] = b;
		candy[b] += candy[a];
		candy[a] = 0;
	}
}

void input() {
	cin >> N >> M >> K;
	for (int i = 1; i <= N; i++) {
		cin >> candy[i];
	}
	for (int i = 0; i < M; i++) {
		int a, b;
		cin >> a >> b;
		friends.push_back({ a,b });
	}
}

void allFriendMerge() {
	init();
	for (pair<int, int> &f : friends) {
		merge(f.first, f.second);
	}
}

void calculateCount() {
	for (int i = 1; i <= N; i++) {
		int present = find(i);
		groupCount[present]++;
	}
}

vector<int> getGroupList() {
	vector<int> result;
	result.push_back(0);
	for (int i = 1; i <= N; i++) {
		if (groupCount[i] > 0) {
			result.push_back(i);
		}
	}
	return result;
}

int findMaxCandy() {
	int result = 0;
	vector<int> groupList = getGroupList();
	int size = groupList.size();
	for (int i = 1; i < size; i++) {
		for (int j = 0; j < K; j++) {
			int present = groupList[i];
			if (j - groupCount[present] >= 0) {
				dp[i][j] = max(dp[i][j], dp[i - 1][j - groupCount[present]] + candy[present]);
			}
			dp[i][j] = max(dp[i][j], dp[i - 1][j]);
			result = max(result, dp[i][j]);
		}
	}
	return result;
}

int solve() {
	allFriendMerge();
	calculateCount();
	return findMaxCandy();
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	input();
	cout << solve() << '\n';
	return 0;
}