백준 21924번 도시 건설 C++

출처 : https://www.acmicpc.net/problem/21924

 

유니온 파인드를 이용해서 최소 스패닝 트리를 만들었다.

 

#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;

struct Info {
    int a;
    int b;
    int cost;
    bool operator<(const Info& other) {
        return cost < other.cost;
    }
};

vector<int> parent;
vector<int> height;
vector<Info> Infos;

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

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

    if (u == v)
        return;

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

    parent[u] = parent[v];

    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);
    for (int i = 1; i <= n; i++) parent[i] = i;

    int64 totalSum = 0;

    for (int i = 0; i < m; i++) {
        int a, b, cost;
        cin >> a >> b >> cost;
        Infos.push_back({ a,b,cost });
        totalSum += cost;
    }
    sort(Infos.begin(), Infos.end());

    int64 sum = 0;
    int cnt = 0;
    for (int i = 0; i < Infos.size(); i++) {
        if (GetParent(Infos[i].a) == GetParent(Infos[i].b))
            continue;
        Merge(Infos[i].a, Infos[i].b);
        sum += Infos[i].cost;
        cnt++;
    }
    int64 ans = totalSum - sum;
    if (cnt < n - 1) ans = -1;
    cout << ans;
}