토마토

• 토마토 문제가 3차원으로 확장된 문제. 동일한 방법으로 BFS를 돌려서 풀면 된다. 해설은 그쪽을 참조.
• 시간 복잡도는 똑같이 O(셀의 갯수) = O(MNH)

"""Solution code for "BOJ 7576. 토마토".

- Problem link: https://www.acmicpc.net/problem/7576
- Solution link: http://www.teferi.net/ps/problems/boj/7576

Tags: [BFS]
"""

from teflib import search

RIPE = '1'
UNRIPE = '0'
EMPTY = '-1'


def main():
    def adj_tomatoes(cur_pos):
        p, x = divmod(cur_pos, M)
        z, y = divmod(p, N)
        delta = 1
        for coord_i, size_i in ((x, M), (y, N), (z, H)):
            if coord_i > 0:
                next_pos = cur_pos - delta
                if tomatoes[next_pos] == UNRIPE:
                    yield next_pos
            if coord_i < size_i - 1:
                next_pos = cur_pos + delta
                if tomatoes[next_pos] == UNRIPE:
                    yield next_pos
            delta *= size_i

    M, N, H = [int(x) for x in input().split()]
    tomatoes = []
    for _ in range(N * H):
        tomatoes.extend(input().split())

    ripe_tomatoes = [i for i, tom in enumerate(tomatoes) if tom == RIPE]
    tomato_count = sum(1 for tom in tomatoes if tom != EMPTY)
    dists = search.min_distances(adj_tomatoes, ripe_tomatoes)

    print(-1 if len(dists) < tomato_count else max(dists.values()))


if __name__ == '__main__':
    main()

• Dependency: teflib.search.min_distances