Python - 토마토 (7569) BFS, 3차원 배열

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

7569번: 토마토

 

3차원 배열은 [z 번째 2차원 배열][x 행][y 열]로 좌표값만 미리 잘 정해놓으면 일반적인 BFS와 거의 똑같았습니다.

import sys
from collections import deque
input = sys.stdin.readline
m, n, h = map(int, input().strip().split())
allSum = m * n * h
ans  = 0
que = deque()
tomato = []
freq = [[[0 for i in range(101)] for j in range(101)] for k in range(101)]
dire = [[-1, 1, 0, 0, 0, 0], [0, 0, -1, 1, 0, 0], [0, 0, 0, 0, -1, 1]]
for i in range(h):
    tomato.append([])
    for j in range(n):
        tomato[i].append(list(map(int, input().strip().split())))
        
        # 음수면 양수로 바꿉니다.
        allSum -= sum(map(lambda x : x if x > 0 else -x, tomato[i][j]))
        
        #토마토 위치를 큐에 초기값으로 설정
        for k in range(m):
            if tomato[i][j][k] == 1:
                que.append((i, j, k))
                freq[i][j][k] = 1
                
if not allSum: 모두 익었다면 0을 출력하고 종료
    print(0)
    exit()

while que:
    if not allSum: # 토마토가 모두 익었다면 출력 후 종료
        print(ans)
        exit()
        
    tmp = que.popleft()
    tz = tmp[0]
    tx = tmp[1]
    ty = tmp[2]
    ans = tomato[tz][tx][ty]
    
    for i in range(6):
        az = tz + dire[2][i]
        ay = ty + dire[1][i]
        ax = tx + dire[0][i]
        
        if az < 0 or az >= h: continue
        if ax < 0 or ax >= n: continue
        if ay < 0 or ay >= m: continue
        if freq[az][ax][ay] or tomato[az][ax][ay] == -1 : continue
        allSum -= 1
        freq[az][ax][ay] = 1
        tomato[az][ax][ay] = tomato[tz][tx][ty] + 1
        que.append((az, ax, ay))
print(-1)