[Python][프로그래머스] lv3. 섬 연결하기

https://school.programmers.co.kr/learn/courses/30/lessons/42861?language=python3 

프로그래머스

 

내 코드

def solution(n, costs):
    answer =0
    parent =[i for i in range(n)]
    costs.sort(key = lambda x:x[2])
    
    def sameParent(parent, a,b):
        return getParent(parent, a) == getParent(parent, b)

    def getParent(parent, x):
        if parent[x] == x:
            return x
        return getParent(parent, parent[x])

    def unionParent(parent, a,b):
        a=getParent(parent, a)
        b = getParent(parent, b)
        if a<b:
            parent[b] =a
        else:
            parent[a]=b
            
    for a,b,cost in costs:
        if sameParent(parent, a,b) == False:
            unionParent(parent, a,b)
            answer += cost
    
    return answer

 

코드 리뷰

크루스칼(kruskal) 알고리즘

크루스칼 알고리즘을 실행할 때 유니온파인드 알고리즘이 정말 흔하게 사용되는데, 그 알고리즘에 대한 개념을 다지는 문제였다고 생각한다. 그래서 lv3이길래 조금 겁먹었는데, 막상 풀어보니 그저 알고리즘 사용 때문에 레벨이 높게 책정된 것이 아닐까 싶다. 

 

문제는 간단하게 다음의 표와 같이 전개된다. 비용 순으로 정렬했던 점을 이용해서 풀이하면 된다.

처리한 간선	parent	parent 설명
[0,1,1]	[0,0,2,3]	0번 노드, 1번 노드가 같은 집합, 1번 노드의 부모가 0이 됨
[1,3,1]	[0,0,2,0]	1번 노드와 3번 노드가 같은 집합, 3번 노드의 부모가 0이 됨
[1,2,5]	[0,0,0,0]	1번 노드와 2번 노드가 같음, 2번 노드의 부모가 0이 됨

 

    answer =0
    parent =[i for i in range(n+1)]
    costs.sort(key = lambda x:x[2])

먼저 기본 세팅을 해준다. 먼저 parent 리스트는 각각의 노드들의 부모를 설명하는 리스트이다. 처음에는 자기 스스로를 부모로 지정한다. 그리고 입력받은 costs를 비용 순으로 정렬해준다. 가장 낮은 비용을 우선으로 생각하기 위해서 정렬을 먼저 해주는 것이 좋다. 

 

    def getParent(parent, x):
        if parent[x] == x:
            return x
        return getParent(parent, parent[x])

유니온파인드 알고리즘 함수들을 그대로 사용해준다. 여기서 parent[x]는 x의 부모를 의미한다. 그래서 x의 부모가 x라면 그대로 x를 반환한다. 이후 union 함수에서 여러 노드들이 들어갈 건데, 해당 노드의 부모 노드를 찾을 때까지 재귀로 계속 타고 들어가서 부모 노드를 찾아낸다.

 

    def unionParent(parent, a,b):
        a=getParent(parent, a)
        b = getParent(parent, b)
        if a<b:
            parent[b] =a
        else:
            parent[a]=b

union함수에서는 비교할 두 노드를 집어넣는다. 그래서 a의 부모 노드를 찾고, b의 부모 노드를 찾는다. union-find 알고리즘에서는 더 작은 부모 노드가 그 집합의 대표자가 되기로 정의했다. 그래서 a와 b를 비교하고, a가 더 작으면 b의 부모를 a로 결정한다.

 

    def sameParent(parent, a,b):
        return getParent(parent, a) == getParent(parent, b)

그리고 same 함수를 통해 a와 b의 부모 노드를 비교한다. 이건 사이클을 만들지 않기 위해 사용되는 중요한 함수이다.  만약 a와 b가 같은 집합 안에 있다면, 이를 연결했을 때 사이클이 발생할 수 있다. 이를 체크하기 위해 이처럼 간단하게 표현할 수 있다. 

 

    for a,b,cost in costs:
        if sameParent(parent, a,b) == False:
            unionParent(parent, a,b)
            answer += cost

마지막으로 입력받은 costs 안에서 각각의 노드와 비용을 꺼내서, same함수를 통해 먼저 사이클이 발생할 수 있는지를 확인한다. 사이클이 발생하지 않는다는 점이 확인되면, union에 넣어서 노드의 부모를 찾는다. 그리고 모든 노드가 이어져야 하는 것이 요구사항이므로, 사이클이 발생하지 않는 선에서 cost를 누적해서 더한다.