[ 백준 ] 15789번 - CTP 왕국은 한솔 왕국을 이길 수 있을까? (C++)

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

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

 

 

[ 문제풀이 ]

 

1. 부모 노드를 저장할 vect 배열을 선언합니다.

 

2. vect[ i ] = i 로 초기화합니다.

 

3. Union-Find를 활용하여 왕국을 이어주고, 이을 때마다 power 배열에 값을 경신해 줍니다.

 

4. for문을 통해 1부터 N까지 돌면서 CTP 왕국이나 한솔 왕국에 속하지 않은 왕국들의 동맹국들의 수를 ans 배열에 저장합니다.

 

5. ans 배열을 내림차순으로 정렬하고, K개까지 동맹국을 더해줍니다.

 

6. 그 결과를 출력합니다.

 

[ 소스코드 ]

 

#include<iostream>
#include<vector>
#include<algorithm>
 
using namespace std;
 
int N, M;
int vect[100001];
int power[100001];
int visited[100001];
vector<int> ans;
 
int Find(int num)
{
    if (vect[num] == num) return num;
    return vect[num] = Find(vect[num]);
}
 
void Union(int a, int b)
{
    int pa = Find(a);
    int pb = Find(b);
 
    vect[pb] = pa;
    power[pa] += power[pb];
}
 
int main()
{
    scanf("%d %d", &N, &M);
 
    for (int i = 1; i <= N; i++) {
        vect[i] = i;
        power[i] = 1;
    }
 
    for (int i = 0; i < M; i++) {
        int X, Y;
        scanf("%d %d", &X, &Y);
 
        if (Find(X) != Find(Y)) {
            Union(X, Y);
        }
    }
 
    int C, H, K;
    scanf("%d %d %d", &C, &H, &K);
    int cnt = 0;
 
    if (K) {
        for (int i = 1; i <= N; i++) {
            int temp = Find(i);
            if ((Find(C) != temp) && (temp != Find(H))) {
                if (visited[temp] == 0) {
                    visited[temp] = 1;
                    ans.push_back({ power[temp] });
                }
            }
        }
    }
 
    sort(ans.begin(), ans.end(), greater<>());
 
    int temp = power[Find(C)];
 
    for (int i = 0; i < min((int)ans.size(), K); i++) {
        temp += ans[i];
    }
    printf("%d", temp);
}