백준 15789번 CTP 왕국은 한솔 왕국을 이길 수 있을까? C++

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

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

 

유니온 파인드

 

#include <iostream>
#include <string>
#include <vector>
#include <math.h>
#include <algorithm>
#include <utility>
#include <stack>
#include <queue>
#include <math.h>
#include <set>
#include <map>
#include <list>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <limits.h>

using namespace std;
using int64 = long long;

int n, m;

vector<int> parent;
vector<int> height;
vector<pair<int, int>> alliances;

int GetParent(int n) {
    if (n == parent[n]) return n;
    return parent[n] = GetParent(parent[n]);
}

void Merge(int u, int v, bool flag = true) {
    u = GetParent(u);
    v = GetParent(v);

    if (u == v) return;

    if (height[u] > height[v])
        ::swap(u, v);

    parent[u] = v;
    if (flag)
        alliances[v].first += alliances[u].first;

    if (height[u] == height[v])
        height[v]++;
}

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

    cin >> n >> m;
    parent = vector<int>(n + 1);
    height = vector<int>(n + 1, 1);
    alliances = vector<pair<int, int>>(n + 1, { 0,0 });
    for (int i = 1; i <= n; i++) {
        parent[i] = i;
        alliances[i] = { 1,i };
    }


    for (int i = 0; i < m; i++) {
        int u, v;
        cin >> u >> v;
        Merge(u, v);
    }
    int c, h, k;
    cin >> c >> h >> k;

    int cnt = 0;

    int ctp = GetParent(c);
    int hansol = GetParent(h);

    int ans = alliances[ctp].first;

    sort(alliances.begin(), alliances.end(), greater<>());

    for (int i = 0; i < n; i++) {
        if (cnt >= k) break;
        int other = GetParent(alliances[i].second);
        int alliance = alliances[i].first;
        if (other == hansol) continue;
        if (other == ctp) continue;
        cnt++;
        ans += alliance;
        Merge(ctp, other, false);
        ctp = GetParent(ctp);
    }
    cout << ans;
}