[Python/파이썬] 백준 18513번 샘터

18513번: 샘터

 

코드

from collections import deque

n, k = map(int, input().split())
waters = list(map(int, input().split()))

q = deque()
visited = set()  # 이미 방문한 위치
for w in waters:
    visited.add(w)  # 샘터 방문 처리
    q.append((w, 0))

houses_num = 0  # 지은 집의 개수
ans = 0  # 불행도의 합
while q:
    now, dist = q.popleft()
    for nx in [now-1, now+1]:
        if nx not in visited:
            q.append((nx, dist+1))
            visited.add(nx)
            houses_num += 1
            ans += (dist+1)
        if houses_num == k:  # k개의 집 건설 완료
            break
    if houses_num == k:
        break
print(ans)

 

BFS는 같은 레벨에 있는 노드를 모두 탐색한 후에 다음 노드를 탐색하는 식으로 진행된다.

이 문제에서는 불행도, 즉 가장 가까운 샘터로부터의 거리가 바로 그 레벨이라고 볼 수 있다. 그러므로 BFS로 탐색하며 탐색된 위치 중 선착순 k개의 위치를 뽑아내면 불행도가 최소가 되도록 집을 지을 수 있다. 

 

처음에는 visited를 평소처럼 10,000,000*2 크기의 배열로 만들었는데 메모리 초과가 떠서 집합으로 바꾸었다. 생각해보면 어차피 BFS는 k개의 위치를 찾을 때까지만 하므로 배열보다는 집합으로 하는 것이 메모리를 효율적으로 사용할 수 있는 방법이다.