백준 1368번: 물대기 (JAVA)

문제

선주는 자신이 운영하는 N개의 논에 물을 대려고 한다. 물을 대는 방법은 두 가지가 있는데 하나는 직접 논에 우물을 파는 것이고 다른 하나는 이미 물을 대고 있는 다른 논으로부터 물을 끌어오는 법이다.

각각의 논에 대해 우물을 파는 비용과 논들 사이에 물을 끌어오는 비용들이 주어졌을 때 최소의 비용으로 모든 논에 물을 대는 것이 문제이다.

입력

첫 줄에는 논의 수 N(1 ≤ N ≤ 300)이 주어진다. 다음 N개의 줄에는 i번째 논에 우물을 팔 때 드는 비용 Wi(1 ≤ Wi ≤ 100,000)가 순서대로 들어온다. 다음 N개의 줄에 대해서는 각 줄에 N개의 수가 들어오는데 이는 i번째 논과 j번째 논을 연결하는데 드는 비용 Pi,j(1 ≤ Pi,j ≤ 100,000, Pi,j = Pj,i, Pi,i = 0)를 의미한다.

출력

첫 줄에 모든 논에 물을 대는데 필요한 최소비용을 출력한다.

 

---

 

내 시도

예제가 저렇다보니, '우물의 비용이 가장 작은 곳에 우물을 대고 나머지는 MST 구하면 되지 않을까?'라고 생각했다. 

반례)

2

1

1

0 2

2 0

이러면 각자 우물을 파서 쓰는게 좋다. 약간 인문학적 감성으로 '우물을 파는 거 보다는 연결하는게 무조건 싸지~'라는 생각때문에 이런 풀이가 나온 듯하다.

 

풀이

우물을 새로운 노드로 보고 MST 구하면 된다.

이런 식으로. (출처: https://blog.encrypted.gg/1024?category=773649)

 

코드

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

public class Main {
    static int[] parents;

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

        int[] prices = new int[N];
        for (int i = 0; i < N; i++) {
            prices[i] = Integer.parseInt(br.readLine());
        }
        Arrays.sort(prices);

        int[][] non = new int[N][N];
        for (int i = 0; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            for (int j = 0; j < N; j++) {
                non[i][j] = Integer.parseInt(st.nextToken());
            }
        }

        PriorityQueue<Edge> pq = new PriorityQueue<>();
        int ans = prices[0];

        for (int i = 0; i < N; i++) {
            for (int j = i+1; j < N; j++) {
                pq.offer(new Edge(i, j, non[i][j]));
            }
        }

        parents = new int[N];
        for (int j = 0; j < N; j++) {
            parents[j] = j;
        }

        // 연결이 다 있다는 보장이 있나? 문제 조건에 보니까 있다.

        int count = 0;
        while (count < N-1) {
            Edge edge = pq.poll();
            if (!union(edge.from, edge.to)) continue;
            ans += edge.weight;
            count++;
        }

        System.out.println(ans);
    }


    static class Edge implements Comparable<Edge>{
        int from;
        int to;
        int weight;

        public Edge(int from, int to, int weight) {
            super();
            this.from = from;
            this.to = to;
            this.weight = weight;
        }

        @Override
        public int compareTo(Edge o) {
            return Integer.compare(this.weight, o.weight);
        }
    }


    static int findRoot(int a) {
        if (parents[a] == a) return a;
        return parents[a] = findRoot(parents[a]);
    }


    static boolean union(int a, int b) {
        int rootA = findRoot(a);
        int rootB = findRoot(b);
        if (rootA == rootB) return false;
        parents[rootB] = rootA;
        return true;
    }
}