MST 문제로 간선 비용이 특정 수치보다 높을 때에만 우선순위 큐에 넣고 크루스칼 알고리즘을 사용하면 되는 문제. 시간/메모리 관련은 (적어도 파이썬에서는) 간선의 개수가 정점의 개수
n
보다 1 적은n-1
일 때 바로 탈출하면 된다.
import sys
import heapq
n, base_cost = map(int, sys.stdin.readline().rstrip().split())
nodes = []
for _ in range(n): nodes.append(list(map(int, sys.stdin.readline().rstrip().split())))
parents = [i for i in range(n)]
pq = []
for i in range(n):
for j in range(i+1, n):
a, b = nodes[i]
c, d = nodes[j]
cost = abs(a-c)**2 + abs(b-d)**2
if cost >= base_cost: heapq.heappush(pq, [cost, i, j])
def find(node):
if parents[node] == node: return node
else:
parents[node] = find(parents[node])
return parents[node]
def union(node1, node2):
root1, root2 = find(node1), find(node2)
if root1 == root2: return False
else:
parents[root2] = root1
return True
total = 0
edge_num = 0
while pq:
cur_cost, node1, node2 = heapq.heappop(pq)
if union(node1, node2):
total += cur_cost
edge_num += 1
if edge_num == n-1: break
if edge_num == n-1: print(total)
else: print(-1)