코딩 테스트/문제 풀기

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

hazel_ 2021. 4. 30. 17:14

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):
    parents=[0]*n
    for i in range(n):
        parents[i]=i

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

    def find_parent(x):
        if parents[x]!=x:
            parents[x]=find_parent(parents[x])


        return parents[x]

    def union(a,b):
        p_a=find_parent(a)
        p_b=find_parent(b)

        if p_a==p_b:
            parents[a]=p_a
            parents[b]=p_a
        elif p_a < p_b:
            if parents.count(p_b)>1:
                while p_b in parents:
                    parents[parents.index(p_b)]=p_a
            parents[p_b]=p_a
        else:
            if parents.count(p_a)>1:
                while p_a in parents:
                    parents[parents.index(p_a)]=p_b
            parents[p_a]=p_b

    for i in range(len(costs)):
        if parents[costs[i][0]] != parents[costs[i][1]]:
            total_cost+=costs[i][2]
            union(costs[i][0], costs[i][1])

    return total_cost