BOJ (Java) 17490번: 일감호에 다리 놓기

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

---

풀이

노드끼리의 연결 정보를 저장한 후, 이를 바탕으로 각 그룹으로 묶어준다.

이 때 모두 하나의 그룹에 묶인다면 와우도를 사용하지 않아도 되므로 YES 출력 후 프로그램을 종료한다.

1개 이상이라면 와우도까지의 돌의 개수를 오름차순으로 정렬, 가장 돌이 적게 드는 곳부터 연결한다.

이 때 해당 그룹에서 다리가 이미 연결되었다면 다른 다리가 추가로 필요하지 않으므로 set에 그룹의 정보를 저장한 후, 한 그룹에서 하나의 다리만 연결되게 한다.

연결 도중 돌의 개수가 K개가 넘었다면 NO를 출력하고, 아니라면 YES를 출력한다.

---

Java 코드

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

public class Main {
    static int N, M;
    static long K;
    static int[] groups;
    static boolean[] isConnected;
    static List<int[]> stones = new ArrayList<>();
    public static void main(String args[]) throws Exception {
        BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));
        String[] inputs = bf.readLine().split(" ");
        N = Integer.parseInt(inputs[0]);
        M = Integer.parseInt(inputs[1]);
        K = Long.parseLong(inputs[2]);

        // save stone info
        inputs = bf.readLine().split(" ");
        for (int i = 1; i <= N; i++) {
            stones.add(new int[] {Integer.parseInt(inputs[i - 1]), i});
        }

        isConnected = new boolean[N + 1];
        Arrays.fill(isConnected, true);
        for (int i = 0; i < M; i++) {
            inputs = bf.readLine().split(" ");
            int from = Integer.parseInt(inputs[0]);
            int to = Integer.parseInt(inputs[1]);

            if (from < to || from == N)
                isConnected[from] = false;
            else
                isConnected[to] = false;
        }

        // group init
        groups = new int[N + 1];
        for (int i = 1; i <= N; i++) {
            groups[i] = i;
        }

        // connect groups
        for (int i = 1; i < N; i++) {
            if (!isConnected[i]) continue;
            union(i, i + 1);
        }
        if (isConnected[N]) {
            union(N, 1);
        }

        // 하나의 그룹에 묶여 있는지 확인
        Set<Integer> s = new HashSet<>();
        for (int i = 1; i <= N; i++) {
            find(i);
            s.add(groups[i]);
        }
        if (s.size() == 1) {
            System.out.println("YES");
            return;
        }
        s.clear();
        
        stones.sort((a, b) -> Integer.compare(a[0], b[0]));
        long stoneCnt = 0;
        for (int[] stone : stones) {
            int cnt = stone[0];
            int node = stone[1];

            // 해당 그룹에서 이미 와우도에 연결했다면 제외
            if (s.contains(groups[node])) continue;

            stoneCnt += cnt;
            if (K < stoneCnt) {
                System.out.println("NO");
                return;
            }
            
            // 해당 그룹 추가
            s.add(groups[node]);
        }

        System.out.println("YES");
    }

    static void union(int from, int to) {
        groups[find(from)] = find(to);
    }

    static int find(int node) {
        if (node == groups[node]) return node;
        return groups[node] = find(groups[node]);
    }
}