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

프로그래머쓱

해당 글에서는 섬 연결하기 문제를 Python을 이용해 풀이하고자 한다.

프로그래머스

 

🔷 문제 설명

 

🔷 문제 풀이

사실 누가 봐도 최소 신장 알고리즘으로 풀이하는 문제이다. n개의 섬이 모두 연결되도록 다리가 연결이 돼야 하며, 최소 비용이 돼야하기 때문이다.

하지만 이 문제를 처음 마주했을 때, 최소 신장 알고리즘을 알지도 못했다. 분명히 학부 수업에서 배웠으나 기억조차 나지 않았다.

 

따라서 다시 공부하는 시간을 가졌다. 혹시나 크루스칼 알고리즘을 모른다면 아래의 글을 참고하도록 하자.

[알고리즘] 크루스칼 알고리즘이란? (최소 신장 트리 찾기)

 

처음엔 아무것도 모르고 다익스트라로 모든 노드 시작점에 대해 거리 테이블을 구하고, 조합을 이용해 모든 경우의 수를 완탐을 통해 최소값을 확인하면 되겠다! 라고 생각했는데 틀린 부분이 있었다. 그래서 크루스칼 알고리즘을 공부하고 이 글까지 쓰게 됐다.

 

아래는 잘 못 푼 처음 코드... (악 더러워!)

# 다익스트라 / 플로이드
import heapq
from itertools import permutations
def solution(n, costs):
    
    INF = int(1e9)
    answer = INF

    graph = [[] for _ in range(n)]
    for a, b, cost in costs:
        graph[a].append((b, cost))
        graph[b].append((a, cost))
    
    distance = []
    for s in range(n):
        distance.append(dijkstra(s, graph))

    per = permutations([i for i in range(n)])
    for p in per:
        cnt = 0
        for i in range(n-1):
            first, second = p[i], p[i+1]
            cnt += distance[first][second]
        answer = min(answer, cnt)
        
    return answer

def dijkstra(start, graph):
    
    INF = int(1e9)
    distance = [INF] * len(graph)
    distance[start] = 0
    
    q = []
    heapq.heappush(q, (0, start))
    while q:
        cost, now = heapq.heappop(q)
        if distance[now] < cost:
            continue
        
        for g in graph[now]:
            d = cost + g[1]
            if d < distance[g[0]]:
                distance[g[0]] = d
                heapq.heappush(q, (d, g[0]))
    
    return distance

 

1️⃣ Python 풀이

def find(parent, a):
    if parent[a] != a: # 부모가 루트 노드가 아닌 경우
        parent[a] = find(parent, parent[a])
    return parent[a]

def union(parent, a, b):
    a = find(parent, a)
    b = find(parent, b)
    
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

def solution(n, costs):
    answer = 0
    
    graph = []
    parent = [i for i in range(n+1)]
    for a, b, cost in costs:
        graph.append((cost, a, b))
        
    graph.sort() # cost를 기준으로 정렬
    for cost, a, b in graph:
        if find(parent, a) != find(parent, b): # 사이클이 생기지 않으면
            union(parent, a, b) # 합치기
            answer += cost # 최소 비용 더하기
    
    return answer