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

'연결할 수 없는 섬은 주어지지 않습니다.' 라는 조건이 있기 때문에

최단 거리를 구하는 대표적인 알고리즘인 프림 알고리즘과 크루스칼 알고리즘을 사용하여 풀 수 있다.

프림 알고리즘 (내 풀이)

import math


def solution(n, costs):
    connect = {costs[0][0]}
    answer = 0

    while len(connect) < n:
        idx = get_min_edge_idx_to_connected_nodes(costs, connect)

        connect.update([costs[idx][0], costs[idx][1]])
        answer += costs[idx][2]
        costs.pop(idx)

    return answer


def get_min_edge_idx_to_connected_nodes(costs, connect):
    min_edge = math.inf
    idx = 0

    for i in range(len(costs)):
        if costs[i][0] in connect and costs[i][1] in connect:
            continue

        if costs[i][0] in connect or costs[i][1] in connect:
            if costs[i][2] < min_edge:
                min_edge = costs[i][2]
                idx = i

    return idx

크루스칼 알고리즘 (다른 사람의 풀이)

parent = {}
rank = {}


# 정점을 독립적인 집합으로 만든다.
def make_set(v):
    parent[v] = v
    rank[v] = 0


# 해당 정점의 최상위 정점을 찾는다.
def find(v):
    if parent[v] != v:
        parent[v] = find(parent[v])

    return parent[v]


# 두 정점을 연결한다.
def union(v, u):
    root1 = find(v)
    root2 = find(u)

    if root1 != root2:
        # 짧은 트리의 루트가 긴 트리의 루트를 가리키게 만드는 것이 좋다.
        if rank[root1] > rank[root2]:
            parent[root2] = root1
        else:
            parent[root1] = root2

            if rank[root1] == rank[root2]:
                rank[root2] += 1


def kruskal(n, costs):
    for v in range(n):
        make_set(v)

    mst = []

    edges = [(costs[i][2], costs[i][0], costs[i][1]) for i in range(len(costs))]
    edges.sort()

    for edge in edges:
        weight, v, u = edge

        if find(v) != find(u):
            union(v, u)
            mst.append(edge)

    return mst


def solution(n, costs):
    answer = 0
    temp = kruskal(n, costs)

    for i in range(len(temp)):
        answer += temp[i][0]

    return answer