白駿16236号サメ

白駿16236号サメ

コード#コード#

import sys
from collections import deque
from heapq import heappop, heappush

def feed(start, scaleCount, babyScale):
  global answer 
  r, c = start
  smallFishies = []
  deq = deque()
  deq.append([r, c, 0])
  visited = []
  maxDist = 400
  while deq:
    r, c, count = deq.popleft()
    if count > maxDist:
      continue
    if board[r][c] <  babyScale and board[r][c] != 0:
      heappush(smallFishies, (count, r, c))
      maxDist = count
    if [r, c] not in visited:
      for i in range(4):
        new_r = r + dy[i]
        new_c = c + dx[i]
        if new_r < 0 or new_r >= N or new_c < 0 or new_c >= N:
          continue
        if board[new_r][new_c] <= babyScale:
          deq.append([new_r, new_c, count+1])
      visited.append([r, c])
  
  if len(smallFishies) == 0:
    print(answer)
  else:
    dist, newRow, newColumn = heappop(smallFishies)
    board[newRow][newColumn] = 0
    answer += dist
    scaleCount+=1
    if scaleCount == babyScale:
      scaleCount = 0
      babyScale += 1
    feed([newRow, newColumn], scaleCount, babyScale)

N = int(sys.stdin.readline())
board = [[0 for _ in range(N)] for _ in range(N)]
start = [0, 0]

for i in range(N):
  li = list(map(int, sys.stdin.readline().split()))
  for j in range(N):
    if li[j] == 9:
      start = [i, j]
      li[j] = 0  
    board[i][j] = li[j]

dy = [-1, 0, 1, 0]
dx = [0, 1, 0, -1]

answer = 0
feed(start, 0, 2)

問題を解く

アルゴリズムフロー