코드
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개의 위치를 찾을 때까지만 하므로 배열보다는 집합으로 하는 것이 메모리를 효율적으로 사용할 수 있는 방법이다.
'문제풀이 > DFS_BFS' 카테고리의 다른 글
[Python/파이썬] 백준 15918번 랭퍼든 수열쟁이야!! (1) | 2024.03.23 |
---|---|
[Python/파이썬] 백준 14716번 현수막 (0) | 2024.03.20 |
[Python/파이썬] 백준 2233번 사과나무 (0) | 2024.02.20 |
[Python/파이썬] 백준 17073번 나무 위의 빗물 (0) | 2024.02.19 |
[Python/파이썬] 백준 10819번 차이를 최대로 (0) | 2024.02.08 |