[프로그래머스] lv3. 섬 연결하기 풀이 (Python, JavaScript)

프로그래머스

lv3. 섬 연결하기

접근한 방법

최소 신장 트리를 찾아낼 수 있는지를 묻는 기본 문제입니다. 프림, 크루스칼 알고리즘 중 선택해서 풀이하면 되고 저는 크루스칼을 이용해서 풀이했습니다.

 

크루스칼 알고리즘은 유니온 파인드 자료구조를 이용해서 간단히 구현할 수 있습니다. 파인드(findP)는 노드의 부모 노드를 찾는 함수이며, 유니온(unionP)은 두 노드 집합의 부모를 연결해주는 함수입니다. 

 

주어진 costs 배열을 비용순으로 정렬해주는 과정이 꼭 필요하고 정렬된 엣지를 순회하면서 싸이클이 없는 경우에만 두 노드의 부모를 연결지어주는 과정을 반복하면 최소 신장 트리를 구할 수 있습니다.

 

Python

def solution(n, costs):
    p = [i for i in range(n)]
    def findP(v):
        if p[v] != v:
            return findP(p[v])
        return v
    def unionP(a, b):
        a, b = findP(a), findP(b)
        if a < b:
            p[b] = a
        else:
            p[a] = b
    result = 0
    for cost in sorted(costs, key = lambda x: x[2]):
        a, b, c = cost
        if findP(a) != findP(b):
            unionP(a, b)
            result += c
    return result

JavaScript

function solution(n, costs) {
    const p = Array(n).fill(0)
    let result = 0;
    for (let i = 0; i < n; i += 1) p[i] = i
    costs.sort((a, b) => a[2] - b[2])
    const findP = (v) => {
        if (p[v] !== v) return findP(p[v]);
        return v;        
    }
    const unionP = (a, b) => {
        a = findP(a)
        b = findP(b)
        if (a < b) p[b] = a;
        else p[a] = b;
    }
    for (const cost of costs) {
        const [a, b, c] = cost
        if (findP(a) !== findP(b)) {
            unionP(a, b)
            result += c
        }
    }
    return result
}