백준 16202 - MST 게임

문제

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

리뷰

크루스칼 알고리즘을 이용하여 풀이할 수 있는 문제였다.

주어진 게임의 규칙에 따르면 매 턴마다 그래프에서 가장 가중치가 작은 간선을
제거하고 MST를 형성할 수 있을 경우 MST의 총 비용을 구하고, MST를 형성할 수
없을 경우 0을 출력해야 한다.

답을 도출하기 위해 크루스칼 알고리즘을 응용하였으며, 간선을 표현하기 위해
Edge 라는 클래스를 산정하였다. 원 그래프의 모든 간선은 입력을 받을 시
비용 기준 최소힙(pq)에 저장해놓았고, 크루스칼 알고리즘은 실행시 이 힙을 복사하여
MST를 형성한다. 그리고 MST를 형성할 수 있을 경우에는 MST 총 간선 비용 합을
반환하고 이외의 경우에는 0을 반환토록 구성하였다.

main에서 K번의 턴이 진행되는 동안 한 턴마다 크루스칼 로직에서 복사하여
사용하는 원 그래프의 간선 정보를 저장한 pq 최소 힙에서 간선을 하나 poll()
하는 것으로 규칙도 올바르게 구현하였다.

로직의 시간복잡도는 크루스칼 로직이 실행되는 while문 부분에서 $O(KNlogM)$으로
수렴하며, 이는 $K=100$, $N=1,000$, $M=10,000$인 최악의 경우에도 무난히
제한 조건 2초를 통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

import static java.lang.Integer.parseInt;

public class Main {
    static int N;
    static int[] parent;
    static PriorityQueue<Edge> pq = new PriorityQueue<>(Comparator.comparingInt(e -> e.w));

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = parseInt(st.nextToken());
        int M = parseInt(st.nextToken());
        int K = parseInt(st.nextToken());

        parent = new int[N + 1];

        int u, v;
        for (int w = 1; w <= M; w++) {
            st = new StringTokenizer(br.readLine());
            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());

            pq.offer(new Edge(u, v, w));
        }

        StringBuilder sb = new StringBuilder();
        while (K-- > 0) { // KNlogM
            sb.append(kruskal() + " ");
            pq.poll();
        }

        System.out.print(sb);
        br.close();
    }

    static int kruskal() { // NlogM
        Arrays.fill(parent, -1);
        PriorityQueue<Edge> temp = new PriorityQueue<>(pq); // log N
        int selectedEdges = 0;
        int totalWeight = 0;
        int r1, r2;

        while (!temp.isEmpty() && selectedEdges < N - 1) {
            Edge e = temp.poll();

            r1 = find(e.u);
            r2 = find(e.v);

            if (r1 == r2) continue;

            if (parent[r1] < parent[r2]) {
                parent[r1] += parent[r2];
                parent[r2] = r1;
            } else {
                parent[r2] += parent[r1];
                parent[r1] = r2;
            }

            selectedEdges++;
            totalWeight += e.w;
        }

        if (selectedEdges == N - 1) return totalWeight;
        else return 0;
    }

    static int find(int u) {
        if (parent[u] < 0) return u;

        return parent[u] = find(parent[u]);
    }

    static class Edge {
        int u, v, w;

        public Edge(int u, int v, int w) {
            this.u = u;
            this.v = v;
            this.w = w;
        }
    }
}

결과