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

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

프로그래머스

풀이

최소 신장 트리 (MST)의 Kruskal 알고리즘을 이용하여 풀었음.

 

우선 cost에 따라 오름차순으로 배열 정렬.

vertex의 연결을 저장하고, Union-Find에 이용할 dicionary를 하나 생성하고,

union-find를 구현(find, union func.)

 

Kruskal 알고리즘을 적용해

cost가 낮은 순으로 정렬한 edge들에 대해 순회하며

사이클이 생기지 않도록 간선(edge)를 연결.

(두 vertex를 연결할 때, root 노드가 동일하면 같은 set에 존재하는 것이므로 연결하게 되면 사이클이 된다.)

from collections import deque
def solution(n, costs):
    # sort by cost
    costs = sorted(costs, key=lambda x: x[2])
    costs = deque(costs)

    # vertex
    v = set([i[0] for i in costs])
    v2 = set([i[1] for i in costs])
    v.update(v2)

    # make union-find dictionary
    connect = {i: i for i in v}

    # feat union-find
    def find(x):
        if connect[x] == x:
            return x
        connect[x] = find(connect[x])
        return connect[x]

    def union(x, y):
        rx = find(x)
        ry = find(y)
        if rx < ry:
            connect[ry] = rx
        else:
            connect[rx] = ry

    answer = 0
    for v1, v2, cost in costs:
        # same root vertex = in same set.
        # so, find(v1) == find(v2) : cycle occur.
        if find(v1) != find(v2):
            union(v1, v2)
            answer += cost
            
    return answer