아기 상어(백준 16236, python)

자세한 설명

https://jiwonna52.tistory.com/11
설명은 위에 내가 더 자세히 써 둔 티스토리를 참고하면 좋을 거 같다.

코드

import sys
from collections import deque

n = int(sys.stdin.readline().rstrip())
graph = []

dY = [-1, 0, 1, 0]
dX = [0, 1, 0, -1]

# 초기 그래프 입력
for i in range(n):
    graph.append(list(map(int, sys.stdin.readline().split(" "))))


# 현재 상어의 크기로 먹을 수 있는 물고기까지 가는 모든 거리를 구하는 함수
def findFishes(shark, sharkY, sharkX):
    q = deque()
    q.append((0, sharkY, sharkX))
    visit = [[False for _ in range(n)] for _ in range(n)]
    visit[sharkY][sharkX] = True
    count = []

    while q:
        cL ,cY, cX = q.popleft()

        for k in range(4):
            nY = cY + dY[k]
            nX = cX + dX[k]

            if 0 <= nY < n and 0 <= nX < n:
                if graph[nY][nX] <= shark: # 지나갈 수 있는 위치
                    if visit[nY][nX] == False:
                        visit[nY][nX] = True
                        q.append((cL+1, nY, nX)) # 지나갔다는 표시로 큐에 넣어준다.
                        if graph[nY][nX] != 0 and graph[nY][nX] < shark: # 만일 칸에 물고기가 있는데 아기 상어의 크기보다 작아서 먹을 수 있으면 리스트에 넣는다.
                            count.append((cL+1, nY, nX))

    # 먹을 수 있는 물고기가 존재
    if len(count) != 0:
        # 거리가 제일 짧은 것 -> 제일 위에 있는 것 -> 제일 왼쪽에 있는 것
        newCount = sorted(count, key=lambda x: (x[0], x[1], x[2]))

        # print(newCount)
        return newCount

    return count


shark = 2
answer = 0
startY = 0
startX = 0
saveFish = 0

for i in range(n):
    for j in range(n):
        if graph[i][j] == 9:
            startY = i
            startX = j
            graph[i][j] = 0
            break

while True:
    fishDist = []
    fishDist = findFishes(shark, startY, startX)

    # 더이상 먹을 수 있는 물고기가 없으면 엄마 호출 -> 반복문 끝
    if len(fishDist) == 0:
        break

    # 정렬된 리스트에서 값 하나 뽑아옴
    c, y, x = fishDist[0]

    # 정답에 거리 더해서 추가
    answer += c

    # 물고기 하나 먹었음 표시
    saveFish += 1
    # 지금까지 먹은 물고기가 아기 상어의 크기라면 아기 상어의 크기를 하나 올려주고 지금까지 먹은 물고기를 0으로 초기화한다.
    if saveFish == shark:
        shark += 1
        saveFish = 0

    # 물고기를 먹었다는 표시로 0
    graph[y][x] = 0
    # 상어가 물고기를 먹기 위해 움직인 위치
    startY = y
    startX = x
    # print(y, x, shark, answer)


print(answer)