백준1197 - 최소 스패닝 트리 (크루스칼) Python

https://www.acmicpc.net/problem/1197

해당 문제는 그래프가 주어졌을 때, 최소 스패닝 트리(모든 정점이 연결된 트리)의 가중치를 구하는 문제다. 처음에는 다익스트라 알고리즘이나 플로이드 워셸 알고리즘으로 접근했다가 도저히 답을 찾지 못해 구글의 도움을 얻었고, 이때 알게 된 크루스칼, Union-Find 알고리즘을 적용하게 되었다.

크루스칼, Find-Union 알고리즘에 대한 설명은 아래 포스팅으로 대체한다.

2023.12.14 - [개발/algorithm] - 크루스칼 알고리즘과 Find-Union 알고리즘 Python

 

크루스칼 알고리즘, Find-Union 알고리즘을 통해, 각 노드들이 하나의 루트 원소를 가지게끔하여 최종적으로 최종적인 최소 비용을 구할 수 있었다.

소스 코드

import sys
data_temp = sys.stdin if sys.platform == 'linux' else open('입력.txt', 'r')
input_data = data_temp.read().splitlines()

def solution (data) :
  arr = [list(map(int,v.split())) for v in data]
  n,_ = arr.pop(0)
  
  root = {}
  for i in range(1, n+1) :
    root[i] = i 
  
  def find(v) :
    if root[v] == v : return v
    else : 
      root[v] = find(root[v])
      return root[v]
    
  def union(a,b) :
    x,y = find(a), find(b)
    root[y] = x

  total = 0
  arr.sort(key= lambda x : x[2])
  for s,e,w in arr : 
    
    if find(s) == find(e) : continue
    total += w
    union(s,e)
    
  print(total)
  
solution(input_data)