Prim’s (MST): Special Subtree

Eunseo·2022년 7월 7일
0

HackerRank

목록 보기
18/18
post-thumbnail

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

✅ Problem Summary

프림 알고리즘(Prim's algorithm)을 이용한 최소 신장 트리(Minimum spanning tree) 탐색 문제


🧮 Applied Theory & Algorithm

1. 프림 알고리즘(Prim's algorithm)

그림 출처: Wikipedia


📑 My Answer

  • 모든 테스트 케이스 통과
from collections import defaultdict
import heapq

def prims(n, edges, start):
    visited = [0] * (n + 1)
    graph = defaultdict(list)
    for u, v, w in edges:
        graph[u].append([w, u, v])
        graph[v].append([w, v, u])

    heapq.heapify(graph[start])
    visited[start] = 1
    total_weight = 0

    while graph[start]:
        w, u, v = heapq.heappop(graph[start])
        if visited[v] == 0:
            visited[v] = 1
            total_weight += w
            for edge in graph[v]:
                if visited[edge[2]] == 0:
                    heapq.heappush(graph[start], edge)

    return total_weight

📌 코드 구현 설명

  • 간선 정보가 저장된 edges를 이용하여 graph를 만든다.
  • heapq 라이브러리를 사용하여 graph최소 힙(Min heap) 형태로 만든다.
    • 가중치가 가장 작은 간선(edge)을 우선적으로 선택해야 하므로 최소 힙을 사용한다.
  • 모든 노드를 방문할 때 까지 graph를 탐색한다.
    • visited에 방문 여부를 기록하여 재방문하지 않도록 한다.
    • 노드를 방문할 때 마다 total_weight에 가중치를 더하고, 반복문이 종료되면 return하여 정답을 구한다.

profile
내가 공부한 것들

0개의 댓글