어두운 길 코드 및 해설 (파이썬)

출처: 이것이 취업을 위한 코딩 테스트다 with 파이썬, Q 43, p 397

 

크루스칼 알고리즘을 이용해 최소 비용의 신장 트리를 탐색했습니다. 

최종 반환해야 할 답이 절약할 수 있는 최대 금액이므로, 전체 도로에 가로등을 설치하는데 드는 비용을 total_cost로 저장하고, 최소 신장 트리를 만들었을 때 가로등을 설치하는데 드는 비용을 min_cost로 저장하여 둘의 차를 반환했습니다.

 

import sys

# 입력 받아오기 
inputs = list(map(int, sys.stdin.read().split()))
# 집의 수 
N = inputs[0]
parent = [0] * N 
for i in range(N):
    parent[i] = i
    
# 도로의 수 
M = inputs[1]

# (도로의 길이, 도로가 잇는 집1, 도로가 잇는 집2)
edges = []
total_cost = 0
for i in range(0, len(inputs[2:]), 3):
    a, b, cost = inputs[2:][i], inputs[2:][i+1], inputs[2:][i+2]
    edges.append((cost, a, b))
    total_cost += cost
    
# 도로 길이가 짧은 순으로 정렬
edges.sort()

def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]

def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b 

# 크루스칼 알고리즘 
min_cost = 0
for edge in edges:
    cost, a, b = edge
    if find_parent(parent, a) != find_parent(parent, b):
        union_parent(parent, a, b)
        min_cost += cost 
        
print(total_cost - min_cost)