[백준] 16562. 친구비(Java)

문제

16562번: 친구비

풀이

bfs, dfs, MST로 푸는 방법이 다양하지만, 복습을 할겸 MST 방법으로 풀어보았습니다.

MST에서는 union-find 알고리즘을 사용했습니다.

 

1. union-find 알고리즘으로 집합을 구성해주고

2. parents 배열을 기반으로 list 배열에 해당 친구에서 발생하는 친구비를 add 했습니다.

    ex) 1번 친구와 같은 집합에 속하는 친구가 1, 2, 4라면 1, 2, 4에 해당하는 친구비를 list[1]에 넣어주는 방식

3. 친구비 계산(1~N 순회)

3-1. list[i]의 size가 0 이라면 continue;

3-2. 해당 친구에게 친구비를 주었다면 continue;

       (해당 친구에게 친구비를 주었는지 판단하기 위해 map 자료구조를 활용했습니다.)

3-3. list[i]를 오름차순 정렬하고 0번째 인덱스의 값(가장 저렴한 친구비) + 현재까지 소모한 친구비 <= K 라면

       ans에 0번째 인덱스의 값(가장 저렴한 친구비)을 더해주었습니다.

       이 과정에서 두 값이 합이 K 보다 클 경우 바로 break! 후 "Oh, no" 출력

3-4. ans 값 출력

코드

import java.io.*;
import java.util.*;

public class Main {
    static int[] parents;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int N = Integer.parseInt(st.nextToken());
        int M = Integer.parseInt(st.nextToken());
        int K = Integer.parseInt(st.nextToken());

        parents = new int[N+1];
        int[] arr = new int[N+1];

        st = new StringTokenizer(br.readLine());
        for(int i = 1 ; i <= N ; i ++){
            parents[i] = i;
            arr[i] = Integer.parseInt(st.nextToken());
        }

        for(int i = 0 ; i < M ; i ++){
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            if(find(a) != find(b)) union(a, b);
        }

        List<Integer>[] list = new ArrayList[N+1];
        for(int i = 1 ; i <= N ; i ++){
            list[i] = new ArrayList<>();
        }

        for(int i = 1 ; i <= N ; i ++){
            list[find(parents[i])].add(arr[i]);
        }

        boolean flag = true;
        int ans = 0;
        Map<Integer, Boolean> map = new LinkedHashMap<>();
        for(int i = 1 ; i <= N ; i ++){
            List cur = list[i];

            if(cur.size() == 0) continue;
            if(map.getOrDefault(i, false)) continue;

            map.put(i, true);
            Collections.sort(cur);

            int target = (int) cur.get(0);
            if(K < target + ans){
                flag = false;
                break;
            }
            ans += target;
        }


        System.out.println(flag ? ans : "Oh no");
    }

    private static void union(int x, int y) {
        x = find(x);
        y = find(y);

        if(x > y) parents[x] = y;
        else parents[y] = x;
    }

    private static int find(int x) {
        if(x == parents[x]) return x;
        return parents[x] = find(parents[x]);
    }
}