[BOJ/Gold 4] 백준 15789 CTP 왕국은 한솔 왕국을 이길 수 있을까?(C++)

문제 링크

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

15789번: CTP 왕국은 한솔 왕국을 이길 수 있을까?

 

문제

CTP 왕국은 정말 깊은 역사를 가지고 있다. 선대 김진서 왕부터 시작하여 전현용 왕을 거쳐 … 마침내 김세진이 CTP 왕국의 왕이되었다. 세진이는 재미없는 개그를 정말 싫어했기 때문에 왕이 되자마자 CTP 왕국에서 가장 재미없는 이한솔을 쫓아냈다. 

화가난 한솔이는 자기의 개그에 유일하게 웃어주던 박정률과 함께 한솔 왕국을 세웠다.

그 이후 33년이 지났다 …………. 

어느새 한솔 왕국은 번창하여 CTP 왕국보다 힘이 쎄졌다. 세진이는 다른 왕국과 동맹을 맺어 CTP 왕국의 힘을 길러 한솔 왕국보다 부흥시키려고 한다.  왕국의 힘이란 동맹국의 수를 의미한다.  (예를 들어 동맹이 없는 나라의 힘은 1이다)

왕국간의 동맹의 법칙은 조금 특별해서 만약에 A왕국과 B왕국이 동맹이고 B왕국과 C왕국이 동맹이라면 A왕국과 C왕국도 동맹이 된다. 

CTP 왕국의 왕 세진이는 최대 K번 다른 왕국과 동맹을 맺을 기회를 갖으며, 현재 동맹관계는 CTP 왕국과 한솔 왕국은 동맹이 아니다. 또한 한솔 왕국과 동맹인 왕국과는 동맹을 맺을 수 없으며 K번의 동맹 맺을 기회를 모두 사용하지 않아도 된다.

각 왕국들의 동맹관계와 CTP 왕국의 번호, 한솔 왕국의 번호가 주어질 때 세진이를 도와 CTP 왕국의 힘의 최댓값을 구하여라. 각 왕국의 번호는 1부터 N까지의 자연수로 나타내어지며, 서로 다른 두 왕국이 같은 번호를 갖는 경우는 없다.

입력

입력의 첫째 줄에 왕국의 수 N(3 ≤ N ≤ 100,000)과 동맹 관계의 수 M(1 ≤ M ≤ 200,000)이 주어진다. 이 후 M개의 줄에 X,Y가 주어진다. 이는 X 왕국과 Y 왕국이 동맹이라는 뜻이다.

입력의 마지막 줄에 CTP 왕국의 번호 C와 한솔 왕국의 번호 H와 추가 동맹의 기회 K(0 ≤ K ≤ 100)가 공백으로 구분되어 주어진다. 

주어지는 입력에서 CTP 왕국과 한솔 왕국은 절대로 동맹이 되지 않게 주어진다.

출력

CTP 왕국의 힘의 최댓값을 출력하라. 

예제 입력 1

10 7
1 2
1 3
2 3
1 4
5 6
8 10
7 9
5 9 1

예제 출력 1

6

예제 입력 2

10 7
1 2
1 3
2 3
1 4
5 6
8 10
7 9
5 1 1

예제 출력 2

4

알고리즘 분류

• 유니온 파인드

풀이

유니온 파인드로 연결하면서 같은 동맹인 나라들의 개수를 각각 구한다.

그리고 C, H가 속해 있지 않은 동맹에 속한 나라들의 개수를 큰 순서대로 K개를 C가 속해 있는 동맹에 속한 나라들의 개수에 더한다.

처음에는 같은 동맹에 속한 나라들의 개수를 BFS로 구하려 했는데 시간 초과가 발생하여 다음과 같이 같은 그룹으로 Union하면서 구하는 방법을 찾아서 적용시켰다.

void Union_Group(int X, int Y) {
	int P_X = Find(X);
	int P_Y = Find(Y);
	if (Group_Size[P_X] >= Group_Size[P_Y]) {
		Parent[P_Y] = P_X;
		Group_Size[P_X] += Group_Size[P_Y];
	}
	else {
		Parent[P_X] = P_Y;
		Group_Size[P_Y] += Group_Size[P_X];
	}
}

코드

#include <iostream>
#include <string>
#include <cmath>
#include <set>
#include <map>
#include <stack>
#include <queue>
#include <vector>
#include <algorithm>
#include <climits>
#define FASTIO cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define MAX 100001
#define LL long long
#define INF 1e9

using namespace std;
int N, M, C, H, K;
int Parent[MAX];
int Group_Size[MAX];
bool visited[MAX];
vector<int> Vec;
int answer;

void Init() {
	for (int i = 0; i < MAX; i++) {
		Parent[i] = i;
		Group_Size[i] = 1;
	}
}

int Find(int X) {	
	if (Parent[X] == X) {
		return X;
	}
	return Parent[X] = Find(Parent[X]);
}

void Union_Group(int X, int Y) {
	int P_X = Find(X);
	int P_Y = Find(Y);
	if (Group_Size[P_X] >= Group_Size[P_Y]) {
		Parent[P_Y] = P_X;
		Group_Size[P_X] += Group_Size[P_Y];
	}
	else {
		Parent[P_X] = P_Y;
		Group_Size[P_Y] += Group_Size[P_X];
	}
}

bool Comp(int A, int B) {
	return (A > B);
}

void Input() {
	Init();
	cin >> N >> M;
	for (int i = 0; i < M; i++) {
		int A, B;
		cin >> A >> B;
		if (Find(A) != Find(B)) {
			Union_Group(A, B);
		}
	}
	cin >> C >> H >> K;
}

void Find_Answer() {
	for (int i = 1; i <= N; i++) {
		int P = Find(i);
		if ((P == Find(C)) || (P == Find(H)) || visited[P]) {
			continue;
		}
		Vec.push_back(Group_Size[P]);
		visited[P] = true;
	}
	sort(Vec.begin(), Vec.end(), Comp);
	answer = Group_Size[C];
	if (Vec.size() >= K) {
		for (int i = 0; i < K; i++) {
			answer += Vec[i];
		}
	}
	else {
		for (int i = 0; i < Vec.size(); i++) {
			answer += Vec[i];
		}
	}
	cout << answer << "\n";
}

int main() {
	FASTIO
	Input();
	Find_Answer();

	return 0;
}