[PS / Python] 백준 #10021 Watering the fields

문제 정보

링크: https://www.acmicpc.net/problem/10021
태그: MST

풀이

MST(최소 스패닝 트리) 문제이다.
크루스칼 알고리즘을 써서 풀 수 있는 문제이다.
모든 지점 사이의 거리를 구해서 C 이상인 것들만 pq에 넣고 while문을 돌면서 pq에서 하나씩 팝한다.
만약 pq에서 나온 정점 2개가 유니온이 된다면 답에 더하고, 되지 않는다면 그 두 정점은 이미 MST에 포함되어 있는 것이므로 넘어간다.
만약 연결된 엣지 개수(=유니온이 성공한 횟수)가 N-1개이면 MST가 존재하므로 답(거리의 합)을 출력하고, N-1보다 작다면 -1을 출력한다.

import sys
import heapq
input = sys.stdin.readline

N, C = map(int, input().rstrip().split())
locs = []
for _ in range(N): locs.append(list(map(int, input().rstrip().split())))
parent = [i for i in range(N)]

prev, ans, cnt = 0, 0, 0

pq = []
idx = 0
for i in range(N-1):
    for j in range(i+1, N):
        a, b = locs[i]
        c, d = locs[j]
        dist = (a-c) ** 2 + (b-d) ** 2
        if dist < C: continue
        heapq.heappush(pq, [dist, i, j])

def _find(a):
    if parent[a] == a: return a
    
    parent[a] = _find(parent[a])
    return parent[a]

def _uni(a, b):
    ra, rb = _find(a), _find(b)

    if ra == rb: return False
    else:
        parent[ra] = rb
        return True

ans, e = 0, 0
while pq:
    d, cur, next = heapq.heappop(pq)

    if _uni(cur, next):
        ans += d; e += 1
        if e == N-1:
            break

if e == N-1:
    print(ans)
else:
    print(-1)
    

+

인공지능 쪽으로 진로를 정한 후, ps하는 언어를 c++에 python으로 바꾸었다.
python을 쓰며 느끼는 점은... 시간 초과, 메모리 초과가 정말 잘 난다는 것이었다.
시간 복잡도를 대충 계산해봤을 때 충분할 것 같은데도 시간초과가 나는 경우도 많았고, python으로 갈아탄지 얼마 되지 않았기 때문에 우선순위 큐 같은 각종 자료구조에도 익숙치 않아서 더 헤맸다.
결국 문제 아이디어를 떠올리기 어렵다기 보다, 요즘은 시간초과와 훨씬 더 많이 싸우는 것 같다.
이 문제도 MST라는 것은 금방 알아보았지만 시간초과가 나서 너무 고생했었다...
결국 로직에는 딱히 문제가 없기 때문에, 코드만 보고는 어디서 시간을 잡아먹는 건지 알 수가 없어서 검색을 해서 다른 분의 코드를 참고하였다.
이렇게 계속 깨지다 보면 파이썬에 익숙해질 수 있을까

알게 된것

• N번 연속적으로 좌표를 입력 받을 때는 다음과 같이 코드를 작성하자

locs = []
for i in range(N): locs.append(list(map(int, input().rstrip().split()))))

• python에서 priority queue는 'heappq'라는 자료구조를 import 해서 쓸 수있다.
  리스트를 하나 선언하고 (ex. pq = [])
  heappq.heappush(pq, [dist, i, j]) 이렇게 푸시하고
  dist, cur, next = heappq.heappop(pq) 이렇게 팝할 수 있다