백준1414 - 불우이웃돕기 (크루스칼) Python

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

 

해당 문제는 최소 스패닝 트리 구현 시 불필요한 가중치의 총합을 구하는 문제다. 때문에 크루스칼 알고리즘으로 접근했다.

크루스칼 알고리즘 관련 포스팅

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

 

문제 풀이 과정

주어진 가중치의 값이 영문으로 주어지기 때문에 숫자값으로 변환하고, 이 과정에서 만약 시작점과 도착점이 같다면 불필요한 가중치값으로 판단했고, 추가적으로 크루스칼 알고리즘 동작 시 시작점과 동작점의 의 root가 같다면 불필요한 가중치로 판단하여 결괏값에 포함시켰다.

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

def solution (data) :
  n = int(data.pop(0))
  arr = []
  result = 0

  for i in range(n) : 
    for j in range(n) : 
      v,w = data[i][j], None
      if v == '0' : continue # 가중치가 0인 경우 패스
      elif 97 <= ord(v) <= 122 : w = ord(v) - 96 # 대문자일 경우
      else : w = ord(v) - 38 # 소문자일 경우
      
      if i == j : result += w # 시작점 == 도착점이라면 불필요한 가중치
      else : arr.append([i+1, j+1, w]) # 그것이 아니라면 배열에 추가
  
  roots = {}
  for i in range(1, n+1) :
    roots[i] = i
  
  def find (node) :
    if roots[node] != node :
      roots[node] = find(roots[node])
    return roots[node]
  
  def union (node1, node2) : 
    root1, root2 = find(node1) , find(node2)
    if node1 < node2 : roots[root2] = root1  
    else : roots[root1] = root2
  
  count = 1 # 모두 연결되어있는지 확인하기 위한 변수
  arr.sort(key=lambda x : x[2]) # 가중치 오름차순 정렬
  for s,e,w in arr :
    if find(s) != find(e) : 
      union(s,e)
      count += 1
    else : result += w
  
  print(result if count == n else -1)
    
solution(input_data)