백준 1368 - 물대기

문제

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

리뷰

크루스칼로 풀이할 수 있는 문제였다.

로직의 발상은 간단하다. 한 정점(논)에서

• 물을 직접 파는 경우
• 다른 정점(논)에 연결하는 경우

이렇게 두 경우를 고려하여 비용을 계산해야 하는데, 이 때 물을
직접 파는 경우를 가상의 정점(아래 로직에선 0)을 하나 설정하여 그 정점과의
간선 관계로 나타내주고, 이 간선들을 크루스칼을 통해 MST를 형성할 때 포함시켜
계산해주면 쉽게 답을 구할 수 있다.

로직의 시간복잡도는 간선의 개수가 $N^2$개 이므로
(행렬로 주어지는 간선의 개수 $N^2$개, 이 중 자기 자신에 대한 간선 $N$개 제외
,가상의 정점과의 간선 $N$개를 합하면 최종적으로 $N^2$개)
크루스칼의 복잡도인 $O(N^2logN^2)$으로 수렴하고 이는 $N=300$인 최악의 경우에도
무난히 제한 조건 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.*;

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));
        N = parseInt(br.readLine());

        parent = new int[N + 1];
        int u, v, w;

        for (v = 1; v <= N; v++) {
            w = parseInt(br.readLine());

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

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

                if (u == v) continue;

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

        int cost = kruskal();
        System.out.println(cost);
        br.close();
    }


    static int kruskal() { // O(N^2 log N^2)
        Arrays.fill(parent, -1);
        int cost = 0;
        int selectedNodes = 0;

        while (selectedNodes < 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;
            selectedNodes++;
        }

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

결과