백준 1368번: 물대기

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

 

N개의 논에 물을 대는 방법은 직접 주거나, 다른 논의 물을 길러오는 것인데 각각의 비용이 주어질 때, 최소한의 비용으롤 모든 논에 물을 댈 때의 비용을 구하는 문제이다.

 

비용의 합이 최소가 되어야 하므로 MST(Minimum Spanning Tree)를 그려야 하며 이를 위해 비용에 대한 프림 알고리즘을 적용해야 한다. 먼저, 비용과 이동하려는 논의 번호를 담는 커스텀 class를 구현한다.

static class Water implements Comparable<Water> {
    
    int target, cost;
    
    public Water(int target, int cost) {
        this.target = target;
        this.cost = cost;
    }
    
    @Override
    public int compareTo(Water o) {
        return this.cost - o.cost;
    }
}

 

 

이후 1번 부터 N번 까지 비용을 기준으로 한 우선순위 큐에 넣고 MST를 생성한다. 최소한의 비용을 가진 객체들을 우선순위 큐에서 꺼내면서 이전 비용과 방문 여부를 비교하고 확인하면서 비용을 더해준다. 그 후, 물을 대지 않은 논에 물을 직접 대거나 현재 논에서 길러오는 비용 중 더 작은 비용의 방법을 우선순위 큐에 넣는다.

Queue<Water> pq = new PriorityQueue<>();
int[] costs = new int[n];
boolean[] visited = new boolean[n];
Arrays.fill(costs, Integer.MAX_VALUE);
for (int i = 0; i < n; i++) {
    pq.add(new Water(i, board[i]));
}
int sum = 0;
while (!pq.isEmpty()) {
    Water current = pq.remove();
    if (current.cost > costs[current.target] || visited[current.target]) {
        continue;
    }
    sum += current.cost;
    visited[current.target] = true;
    boolean isFull = true;
    for (int j = 0; j < n; j++) {
        if (!visited[j]) {
            isFull = false;
            break;
            }
    }
    if (isFull) {
        break;
    }
    for (Water next : sendWaterMap.get(current.target)) {
        if (visited[next.target]) {
            continue;
        }
        int min = Math.min(board[next.target], next.cost);
        if (min > costs[next.target]) {
            continue;
        }
        if (board[next.target] > next.cost) {
            pq.add(new Water(next.target, next.cost));
        } else {
            pq.add(new Water(next.target, board[next.target]));
        }
    }
}

 

 

결과 코드는 다음과 같다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Queue;
import java.util.StringTokenizer;

public class 물대기1368 {
    
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int n = Integer.parseInt(br.readLine());
        int[] board = new int[n];
        int answer = 0;
        for (int i = 0; i < n; i++) {
            board[i] = Integer.parseInt(br.readLine());
            answer += board[i];
        }
        StringTokenizer st;
        Map<Integer, List<Water>> sendWaterMap = new HashMap<>();
        for (int i = 0; i < n; i++) {
            sendWaterMap.put(i, new ArrayList<>());
        }
        for (int i = 0; i < n; i++) {
            st = new StringTokenizer(br.readLine());
            for (int j = 0; j < n; j++) {
                int cost = Integer.parseInt(st.nextToken());
                if (i == j) {
                    continue;
                }
                sendWaterMap.get(i).add(new Water(j, cost));
            }
        }
        br.close();
        Queue<Water> pq = new PriorityQueue<>();
        int[] costs = new int[n];
        boolean[] visited = new boolean[n];
        Arrays.fill(costs, Integer.MAX_VALUE);
        for (int i = 0; i < n; i++) {
            pq.add(new Water(i, board[i]));
        }
        int sum = 0;
        while (!pq.isEmpty()) {
            Water current = pq.remove();
            if (current.cost > costs[current.target] || visited[current.target]) {
                continue;
            }
            sum += current.cost;
            visited[current.target] = true;
            //System.out.print(current.cost + " ");
            boolean isFull = true;
            for (int j = 0; j < n; j++) {
                if (!visited[j]) {
                    isFull = false;
                    break;
                }
            }
            if (isFull) {
                break;
            }
            for (Water next : sendWaterMap.get(current.target)) {
                if (visited[next.target]) {
                    continue;
                }
                int min = Math.min(board[next.target], next.cost);
                if (min > costs[next.target]) {
                    continue;
                }
                if (board[next.target] > next.cost) {
                    pq.add(new Water(next.target, next.cost));
                } else {
                    pq.add(new Water(next.target, board[next.target]));
                }
            }
        }
        answer = Math.min(answer, sum);
        System.out.print(answer);
    }
    
    static class Water implements Comparable<Water> {
        
        int target, cost;
        
        public Water(int target, int cost) {
            this.target = target;
            this.cost = cost;
        }
        
        @Override
        public int compareTo(Water o) {
            return this.cost - o.cost;
        }
    }
}