[Python] 섬 연결하기

💻 문제 정의

n개의 섬 사이에 다리를 건설하는 비용(costs)이 주어질 때, 최소의 비용으로 모든 섬이 서로 통행 가능하도록 만들 때 필요한 최소 비용을 return 하도록 solution을 완성하세요.

다리를 여러 번 건너더라도, 도달할 수만 있으면 통행 가능하다고 봅니다. 예를 들어 A 섬과 B 섬 사이에 다리가 있고, B 섬과 C 섬 사이에 다리가 있으면 A 섬과 C 섬은 서로 통행 가능합니다.

💡 접근 및 설계

Union-Find 알고리즘을 이용하여 접근하였다. 각 노드의 부모 노드가 다르다면, Union하면서 동시에 cost를 더한다.

✏️ 알고리즘 풀이

def find_root(x):
    if x == parent[x]: return x
    parent[x] = find_root(parent[x])
    
    return parent[x]

def union(a, b):
    x = find_root(a)
    y = find_root(b)
    
    if x > y:
        parent[x] = y
    else:
        parent[y] = x
        
def same_root(a, b):
    return find_root(a) == find_root(b)

res = 0
for a, b, cost in costs:
    # union-find
    if not same_root(a, b):
        union(a, b)
        res += cost

두 노드의 부모 노드를 찾아 다르다면, 합치고 비용을 계산한다.

🗒️ 최종 제출 코드

import heapq

def solution(n, costs):
    def find_root(x):
        if x == parent[x]: return x
        parent[x] = find_root(parent[x])
        
        return parent[x]
    
    def union(a, b):
        x = find_root(a)
        y = find_root(b)
        
        if x > y:
            parent[x] = y
        else:
            parent[y] = x
            
    def same_root(a, b):
        return find_root(a) == find_root(b)
        
    
    costs.sort(key=lambda x : x[2])
    
    parent = [i for i in range(n)]
    
    res = 0
    for a, b, cost in costs:
        # union-find
        if not same_root(a, b):
            union(a, b)
            res += cost
            
    return res

---

💭 오늘의 회고

-