Python - 아기 상어 (16236) BFS, 우선순위 큐

https://www.acmicpc.net/problem/16236

16236번: 아기 상어

 

1. 더 이상 먹을 수 있는 물고기가 공간에 없다면 아기 상어는 엄마 상어에게 도움을 요청한다.

• 우선순위 큐가 비어있다면 break

2. 먹을 수 있는 물고기가 1마리라면, 그 물고기를 먹으러 간다.

3. 거리는 아기 상어가 있는 칸에서 물고기가 있는 칸으로 이동할 때, 지나야하는 칸의 개수의 최솟값이다.

 

4. 거리가 가까운 물고기가 많다면, 가장 위에 있는 물고기, 그러한 물고기가 여러마리라면, 가장 왼쪽에 있는 물고기를 먹는다.

• 우선순위 큐에 (깊이, x좌표,  y좌표) 순으로 배치 후 저장합니다.

 

 

import sys
import heapq as hq
from collections import deque
input = sys.stdin.readline
h = []
a = []
n = int(input().strip())
sx, sy = 0, 0
s = 2
for i in range(n):
    a.append(list(map(int, input().strip().split())))
    if 9 in a[i]:
        sx, sy = i, a[i].index(9)
        a[i][sy] = 0

que = deque()
dire = [[0, -1], [0, 1], [-1, 0], [1, 0]]
def BFS(x, y):
    global n, s
    que.clear()
    freq = [[0 for _ in range(n)] for __ in range(n)]

    que.append((x, y, 0))
    freq[x][y] = 1
    while que:
        tx, ty, depth = que.popleft()

        # 먹을 수 있다면 먹음
        if a[tx][ty] != 0 and a[tx][ty] < s:
            hq.heappush(h, (depth, tx, ty))

        for t1, t2 in dire:
            ax, ay = tx + t1, ty + t2
            if not (0 <= ax < n and 0 <= ay < n): continue
            # 방문했거나 상어보다 큰 물고기라면
            if freq[ax][ay] or a[ax][ay] > s : continue
            freq[ax][ay] = 1
            que.append((ax, ay, depth + 1))
cnt = 0
ans = 0
while True:
    BFS(sx, sy)
    # 먹을 수 있는 물고기가 없다면 탈출
    if not h: break
    else:
        # 상어 위치 갱신
        # 힙 우선순위 
        # (얼마나 가까운지, 가장 위에 위치한 물고기, 가장 왼쪽에 위치한 물고기)
        ac, sx, sy = hq.heappop(h)
        h.clear()
        a[sx][sy] = 0 # 먹었으니 0으로 바꿈
        cnt += 1
        ans += ac

        # 크기만큼 먹었다면 상어 크기 + 1
        if cnt == s:
            s += 1
            cnt = 0
print(ans)