[백준 10021] Watering the Fields

Junyoung Park·2022년 3월 29일
0

코딩테스트

목록 보기
326/631
post-thumbnail

1. 문제 설명

Watering the Fields

2. 문제 분석

MST 문제로 간선 비용이 특정 수치보다 높을 때에만 우선순위 큐에 넣고 크루스칼 알고리즘을 사용하면 되는 문제. 시간/메모리 관련은 (적어도 파이썬에서는) 간선의 개수가 정점의 개수 n보다 1 적은 n-1일 때 바로 탈출하면 된다.

3. 나의 풀이

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)
profile
JUST DO IT

0개의 댓글