BOJ 16202 : MST 게임 [C++]

 

16202번: MST 게임

 

문제를 이해하는 데 시간이 좀 걸렸었다.

 

구현해야 할 내용을 정리하자면 아래와 같다.

1. 매 턴마다 MST를 만든다.

2. MST가 완성됐다면 MST의 비용을 출력하고, 해당 MST에서 가장 비용이 작은 간선을 제거한다.

3. MST가 완성되지 않았다면 해당 턴부터 마지막 턴까지 0을 출력한다.

 

문제에서 주어진 대로 계산했을 때 간선 개수는 최대 50만 개다.

그렇기 때문에 제거할 간선에 대한 표시를 bool 배열로 해도 되겠다 생각했고, 

매 턴마다 MST를 만들 때 이미 제거된 간선은 무시하도록 해주었다.

 

또한 최대 100 턴까지 진행해도 '50만(간선 배열) X 100(턴 수) = 5천 만'이기 때문에

간선 배열의 처음부터 끝까지 순회하며 MST를 만들어도 시간제한 안에 동작이 가능하다.

 

 

    for (int i = 1; i < M + 1; i++) {
        std::cin >> X >> Y;
        origin.push_back({i, {X, Y}});
    }

    std::sort(origin.begin(), origin.end());

 

간선을 입력받고, 비용 오름차순으로 정렬해 준다.

 

    bool end = false;
    for (int turn = 0; turn < K; turn++) {
        if (end) {
            std::cout << "0 ";
            continue;
        }

        for (int i = 1; i < N + 1; i++) {
            parents[i] = i;
        }

        int turn_cost = 0;
        int cnt = 0;
        for (int i = 0; i < origin.size(); i++) {
            if (erased[i]) {
                continue;
            }
            int cost = origin[i].first;
            X = origin[i].second.first;
            Y = origin[i].second.second;

            int px = find_parent(X);
            int py = find_parent(Y);

            if (px == py) {
                continue;
            }

            if (px < py) {
                parents[py] = px;
            } else {
                parents[px] = py;
            }
            if (cnt == 0) {
                erased[i] = true;
            }
            turn_cost += cost;
            cnt++;
            if (cnt == N - 1) {
                break;
            }
        }

        if (cnt == N - 1) {
            std::cout << turn_cost << " ";
        } else {
            end = true;
            std::cout << "0 ";
        }
    }

 

매 턴마다 MST를 만들 때 union-find를 사용하므로 각 턴의 시작 부분에서 parents 배열을 초기화해준다.

 

이미 앞선 턴에서 MST를 만들고 제거된 간선일 경우 건너뛴다.

서로 같은 집합이 아니라면(부모가 다르다면) 합쳐주고,

MST 비용에 해당 간선의 비용을 더한다.

이때 해당 간선이 MST의 첫 간선이라면, 제거됐음을 표시하는 erase 배열에 표시해 준다.

(이번 턴에서 MST가 만들어지든 만들어지지 않든 제거한다.)

 

MST가 만들어졌다면 비용을 출력

그렇지 않다면 게임이 진행될 수 없다는 표시로 end를 true로 표시하고 0을 출력한다.

 

 

 

전체 코드

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

#define PII std::pair<int, int>

int N, M, K;
int X, Y;
int parents[1001];
std::vector<std::pair<int, PII>> origin;
bool erased[500001];

int find_parent(int x) {
    if (parents[x] == x) {
        return x;
    } else {
        return parents[x] = find_parent(parents[x]);
    }
}

int main() {
    std::cin >> N >> M >> K;

    for (int i = 1; i < N + 1; i++) {
        parents[i] = i;
    }

    for (int i = 1; i < M + 1; i++) {
        std::cin >> X >> Y;
        origin.push_back({i, {X, Y}});
    }

    std::sort(origin.begin(), origin.end());

    bool end = false;
    for (int turn = 0; turn < K; turn++) {
        if (end) {
            std::cout << "0 ";
            continue;
        }

        for (int i = 1; i < N + 1; i++) {
            parents[i] = i;
        }

        int turn_cost = 0;
        int cnt = 0;
        for (int i = 0; i < origin.size(); i++) {
            if (erased[i]) {
                continue;
            }
            int cost = origin[i].first;
            X = origin[i].second.first;
            Y = origin[i].second.second;

            int px = find_parent(X);
            int py = find_parent(Y);

            if (px == py) {
                continue;
            }

            if (px < py) {
                parents[py] = px;
            } else {
                parents[px] = py;
            }
            if (cnt == 0) {
                erased[i] = true;
            }
            turn_cost += cost;
            cnt++;
            if (cnt == N - 1) {
                break;
            }
        }

        if (cnt == N - 1) {
            std::cout << turn_cost << " ";
        } else {
            end = true;
            std::cout << "0 ";
        }
    }

    return 0;
}