본문 바로가기

PS/프로그래머스

[프로그래머스-탐욕법] 섬 연결하기_파이썬(최소신장트리, 크루스칼 알고리즘)

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

 

코딩테스트 연습 - 섬 연결하기

4 [[0,1,1],[0,2,2],[1,2,5],[1,3,1],[2,3,8]] 4

programmers.co.kr


def solution(n, costs):
    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]
    
    def union(a, b):
        a = find(a)
        b = find(b)
        if a < b:
            parent[b] = a
        else:
            parent[a] = b
        
    parent = {}
    for i in range(n):
        parent[i] = i
        
    edges = sorted(costs, key = lambda x: x[2])
    
    cost = 0
    for a, b, c in edges:
        if find(a) != find(b):
            union(a, b)
            cost += c
    
    return cost

find, union 함수 생성.

parent = {}, 자기 자신을 루트노드로 설정.

간선들을 비용 순으로 정렬.

비용이 작은 간선부터 find 함수를 통해 루트노드가 같은지 검사. 다르다면 union, cost += c

for문이 종료되면 return cost


def solution(n, costs):
    def find(x):
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]
    
    parent = [i for i in range(n)]
    
    edges = sorted(costs, key = lambda x: x[2])
    
    cost = 0
    for a, b, c in edges:
        if find(a) != find(b):
            parent[find(a)] = b
            cost += c
    
    return cost

union 함수를 생성하지 않고 문제를 풀이할 수도 있다.

 

find 함수 생성.

parent = [i for i in range(n)] # 자기 자신을 루트노드로 설정.

간선들을 비용 순으로 정렬.

비용이 작은 간선부터 a, b의 루트노드가 같은지 검사.

다르다면, parent[find(a)] = b # a의 루트노드를 b로 설정.(union 함수의 역할을 대체, a와 b가 바뀌어도 상관 없음)

cost += c

for문이 종료되면 return cost