백준 10423 - 전기가 부족해

문제

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

리뷰

크루스칼을 이용해 풀이할 수 있는 간단한 문제였다.

문제 조건만 보면 구현이 복잡할 것 같지만 아주 간단하다.
우선 발전소가 설치된 $K$개의 도시들을 가상의 도시 0과 비용인 0인 간선을
설정하여 비용 기준 최소힙에 저장해준다.

이후, $M$개의 간선 정보를 입력받아 최소힙에 저장해둔 후 크루스칼 로직을
실행하며 MST를 이루는 간선들의 비용 합을 구하면 된다.

MST가 0을 루트로 하여 발전소가 존재하는 도시들을 연결하고 그 도시들에
최소 비용으로 다른 도시들을 연결하는 구조로 형성되기 때문에, 문제에서 예로 든
발전소가 존재하며 다른 도시들에 연결되지 않는 도시같은 경우도 자연스레 처리된다.

실제 로직에선 유니온 파인드에서 경로 압축, 균형 트리를 이용했기 때문에
0을 루트, 나머지 도시들을 자식으로 하는 MST형태가 된다.

로직의 시간복잡도는 간선의 개수가 $M+K$일 때($K$개의 도시 &0과의 간선)
크루스칼의 $O((M+K)log(M+K))$로 수렴하고 이는 $M+K=100,000+1,000$인
최악의 경우에도 무난히 제한 조건 1초를 통과한다.

코드

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.*;

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());

        int u, v, w;

        st = new StringTokenizer(br.readLine());
        while (K-- > 0) {
            v = parseInt(st.nextToken());
            pq.offer(new Edge(0, v, 0));
        }

        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());
            w = parseInt(st.nextToken());

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

        System.out.println(kruskal());
        br.close();
    }

    static int kruskal() {
        parent = new int[N + 1];
        Arrays.fill(parent, -1);
        int mstEdges = 0;
        int cost = 0;

        while (mstEdges < N) {
            Edge e = pq.poll();
            int r1 = find(e.u);
            int 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;
            }

            cost += e.w;
            mstEdges++;
        }

        return cost;
    }

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

결과