[백준] 14621_나만 안되는 연애 C++

10 Dec 2023 in Study on Coding Test

MST를 이용하는 문제

정답제출코드

#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;

typedef pair<int, pair<int, int>> piii;

vector<piii> Edges;
char UnivType[1001];
int N, M, Parent[1001];

int FindParent(int a)
{
    if (Parent[a] == a)
        return a;
    else
        return Parent[a] = FindParent(Parent[a]);
}

void UnionParent(int a, int b)
{
    a = FindParent(a);
    b = FindParent(b);
 
    if (a < b)
        Parent[a] = b;
    else
        Parent[b] = a;
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    cin >> N >> M;

    for (int i = 1; i <= N; ++i)
    {
        cin >> UnivType[i];
        Parent[i] = i;
    }

    for (int i = 0; i < M; ++i)
    {
        int S, E, C;
        cin >> S >> E >> C; // u <-> v는 d만큼 떨어짐
 
        if (UnivType[S] != UnivType[E])
            Edges.push_back({ C, {S, E}});
    }

    sort(Edges.begin(), Edges.end());

    int ans = 0, cnt = 0;
    for (size_t i = 0; i < Edges.size(); ++i)
    {
        int u = Edges[i].second.first;
        int v = Edges[i].second.second;
        int d = Edges[i].first;
 
        if (FindParent(u) != FindParent(v))
        {
            UnionParent(u, v);
            ans += d;
            if (++cnt == N - 1)
            {
                cout << ans;
                return 0;
            }
        }
    }
    cout << -1;

    return 0;
}

MST를 이용하면 되는 문제이다.

다만, 간선이 주어질 때 남자학교 - 남자학교 / 여자학교 - 여자학교 간 간선이 주어질 수도 있으니

이를 제외시키는 기술이 필요하다. 그 부분이 아래 코드이다.

if (UnivType[S] != UnivType[E])
    Edges.push_back({ C, {S, E}});

그리고 나머지는 크루스칼 알고리즘 템플릿을 따라가면서

연결시킨 학교 개수가 N이 된다면, 즉 간선을 유니온 시킨 것이 N-1이 되면 그때 비용을 출력한다.