[프로그래머스] 섬 연결하기

문제 주소

https://school.programmers.co.kr/learn/courses/30/lessons/42861

풀이 과정

다리를 최소 비용으로 건설해서 모든 섬을 연결하는 문제이다.

1차 시도

처음에는 가장 비용이 적게 드는 다리부터 하나씩 건설하되, 이미 섬끼리 다른 다리로 연결되어 있는 경우에는 건설하지 않도록 로직을 짜려고 했다.

방향에는 문제가 없었지만 이를 구현하는 과정에서 부실한 부분이 있었다.

두 섬이 연결되어 있는지를 확인하기 위해 다음 작업을 수행했다.
A섬과 B섬을 연결하는 다리를 건설한다고 하면, A섬에 연결되어 있던 섬들과 B섬이 연결된 것으로 표시하고, 반대로 B섬과 연결되어 있던 섬들도 A섬과 연결된 것으로 표시했다.

이 코드는 최종 테스트를 통과하지 못했고, 이에 대해 고민하던 중 Kruskal Algorithm에 대해 알아보게 되었다.

이와 별개로 아래 코드는 A섬에 연결된 섬들과 B섬에 연결된 섬들간의 연결을 표현하지 못한것이 문제라는 사실을 테스트 케이스 추가를 통해 확인했다.

import java.util.*;
class Solution {
    public int solution(int n, int[][] costs) {
        int[][] bridge = new int[n][n];
        for(int i = 0; i < n; i++) bridge[i][i] = 1;
        
        return Arrays.stream(costs).sorted(Comparator.comparingInt(c -> c[2])).filter(c -> {
            int i1 = c[0];
            int i2 = c[1];
            if(bridge[i1][i2] == 1) return false;
            for(int i = 0; i < n; i++){
                if(bridge[i1][i] == 1) {
                    bridge[i2][i] = 1;
                    bridge[i][i2] = 1;
                }
                if(bridge[i2][i] == 1) {
                    bridge[i1][i] = 1;
                    bridge[i][i1] = 1;
                }
            }
            return true;
        }).mapToInt(c -> c[2]).sum();
    }
}

위 코드는 B와 1,2가 연결됨을 기록하고, A와 3,4가 연결됨을 기록하지만 1,2,3,4간의 연결은 기록하지 못한다.

최종 풀이(Kruskal Algorithm)

크루스칼 알고리즘은 위 문제를 위해 딱 필요한 알고리즘으로, 그래프의 모든 정점을 포함하면서 간선의 가중치가 최소인 트리를 일컫는 최소 신장 트리를 찾는 알고리즘이다.

먼저 다리를 비용순으로 정렬하는 것은 처음에 생각했던것과 같다. 그리고 다리를 하나씩 추가하되, 다리를 추가함으로써 사이클이 형성된다면 추가하지 않는다.

사이클을 형성하는지를 확인하기 위해서, 두 섬을 연결할 때 연결된 섬 그룹에 있는 섬들중 가장 번호를 작은 섬의 번호를 그룹에 있는 각 섬들에 저장한다.

위 그림을 예시로 들면, A = 5, B = 6이라 하면 A와 B를 연결하기 전에는 1, A, 2에는 각각 가장 작은 수인 1이 저장되고, B, 3, 4에는 각각 가장 작은 수인 3이 저장된다.

만약 저장된 수가 같은 두 섬을 연결한다면, 사이클이 형성된다. 즉, A와 2를 연결한다면 두 섬에는 모두 1이 저장되어 있으므로 사이클이 형성된다.

A와 B에는 다른수가 저장되어 있으므로, 연결했을 때 사이클이 생기지 않는다. 따라서 연결할 수 있고, 연결하고 나면 오른쪽에 있던 섬들에는 모두 1이 저장된다. 구체적으로는, A에는 1이 저장되어 있고 B에는 3이 저장되어 있으므로, 연결 후 3이 저장된 모든 섬들을 찾아서 1로 변경한다. 이렇게 하면 간접적으로 연결된 모든 섬들을 확인하는 것이 가능하다.

최종 코드

import java.util.*;
class Solution {
    public int solution(int n, int[][] costs) {
        int[] bridge = new int[n];
        for(int i = 0; i < n; i++) bridge[i] = i;
        
        return Arrays.stream(costs).sorted(Comparator.comparingInt(c -> c[2])).filter(c -> {
            int i1 = c[0];
            int i2 = c[1];
            if(bridge[i1] == bridge[i2]) return false;
            int max = Math.max(bridge[i1], bridge[i2]);
            int min = Math.min(bridge[i1], bridge[i2]);
            for(int i = 0; i < n; i++){
                if(bridge[i] == max) bridge[i] = min;
            }
            return true;
        }).mapToInt(c -> c[2]).sum();
    }
}