백준 16398번 행성 연결 C++

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

16398번: 행성 연결

 

유니온 파인드 자료구조를이용해서 최소 스패닝 트리를 구한다.

 

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

struct Info {
    int s;
    int e;
    int cost;

    bool operator<(const Info& other) const {
        return cost < other.cost;
    }
};

int n;
vector<vector<int>> board;
vector<Info> v;
vector<int> parent;
vector<int> height;

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

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] = v;

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

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

    int64 ans = 0;
    cin >> n;
    board = vector<vector<int>>(n, vector<int>(n));
    parent = vector<int>(n);
    height = vector<int>(n, 1);
    for (int i = 0; i < n; i++) parent[i] = i;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            cin >> board[i][j];
        }
    }
    for (int i = 1; i < n; i++) {
        for (int j = 0; j < i; j++) {
            v.push_back({ i,j,board[i][j] });
        }
    }
    sort(v.begin(), v.end());
    for (auto a : v) {
        if (GetParent(a.s) == GetParent(a.e)) continue;
        Merge(a.s, a.e);
        ans += a.cost;
    }
    cout << ans;
}