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

1. 아이디어

최소 비용으로 모든 섬을 연결한다 -> 최소 비용 트리로 접근해서 해결한다. 따라서 크루스칼 알고리즘으로 해결하면된다.

섬의 개수가 n이라면 모든 노드를 연결하기 위한 필요한 간선의 개수는 n-1이다. 그리고 최소 값을 구하기 위해선, 최소 간선들을 찾으면 되는데 주의해야할 점은 사이클이 만들어지면 안된다. 사이클이 만들어지면 노드를 다시 방문하는 것을 의미하므로 최소 비용을 구할 수 없다. 사이클의 여부는 유니온 파인드를 이용한다.(i.e. 사이클이 이루어지지 않기 위해서는 처음 방문하는 노드를 찾아야한다.)

2. 코드

uf = []

def find(a):
    if uf[a] < 0: return a
    uf[a] = find(uf[a])
    return uf[a]

def merge(a, b):
    a = find(a)
    b = find(b)
    if a == b: return False
    uf[b] = a
    return True

def solution(n, costs):
    global uf
    uf = [-1 for _ in range(n)]
    costs.sort(key=lambda x: x[-1])
    total, cnt = 0, 0
    for a, b, w in costs:
        if merge(a, b):
            total += w
            cnt += 1
            if cnt == n-1: break
    
    return total