HackerRank / Prim's(MST) : Special SubTree / Python

https://www.hackerrank.com/challenges/primsmstsub/problem?isFullScreen=true

Prim's (MST) : Special Subtree | HackerRank

 

프림 알고리즘은 크루스칼과 달리 연결되어 있는 노드에서 최소 간선을 선택해야한다.

이는 기회다.

무조건 정렬을 하고 시작 노드는 연결되어 있고 끝 노드는 연결되어 있지 않은것만 찾으면 된다.

물론, 이 문제 또한 양방향 간선이기 때문에 인풋 노드 둘 중 하나만 연결되어 있는지 확인하면 된다.

 

def prims(n, edges, start):
    visited = []
    distance = 0
    visited.append(start)
    new_edges = edges
    cicle = 0
    while n > cicle:
        new_edges.sort(key = lambda x:x[2])
        for i in range(len(new_edges)):
            if new_edges[i][0] in visited or new_edges[i][1] in visited: # 둘중 하나는 연결되어 있는 노드이고
                if new_edges[i][0] in visited and new_edges[i][1] in visited: #그렇다고 둘다 연결되어 있어서 둘을 연결하면 사이클이 생기는건 방지하고
                    pass
                else:
                    distance += new_edges[i][2]
                    visited.append(new_edges[i][0]) # 사실 매우 보기 안좋지만
                    visited.append(new_edges[i][1]) # 중복되도 상관없고 귀찮으니 이렇게 하자
                    new_edges.pop(i)                # 언젠가 더 좋은 방법이 생각나면 업데이트 예정
                    break
        cicle+=1          
    return distance