프로그래머스

프로그래머스 섬 연결하기 파이썬

hotchya 2021. 9. 12. 00:50

섬 연결하기

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

문제 단순화

  • n개의 섬이 있고, [섬, 섬, 비용]을 나타내는 다리들이 있다.
  • 모든 섬을 통행할 수 있도록 다리를 만들 때, 최소 비용을 리턴한다.

문제풀이

  • 다리를 짓는 비용이 작은 순서로 정렬한다.
  • 다리를 잇는 섬들이 같은 집합에 속해 있는 지 확인한다.(파인드)
  • 같은 집합에 속하지 않은 경우, 같은 집합으로 만든다.(유니온)
  • 비용을 추가한다.
  • n개의 섬이 이어졌을 때, 발생하는 비용을 리턴한다.

코드

def solution(n, costs):
    answer = 0

    p = [i for i in range(n)]

    def find(x):
        if p[x] != x:
            p[x] = find(p[x])
        return p[x]

    def union(x,y):
        px = find(x)
        py = find(y)
        p[py] = px

    costs.sort(key=lambda x : -x[2])

    while n > 0:
        a,b,c = costs.pop()
        if find(a) != find(b):
            union(a,b)
            answer += c
            n -= 1

    return answer