[백준] 23288 주사위 굴리기 2 with 파이썬

import sys
from collections import deque

input = sys.stdin.readline
N, M, K = map(int, input().split())                             # 행, 열, 주사위 이동 횟수
array = [list(map(int, input().split())) for _ in range(N)]     # 격자 맵
dice_x, dice_y = 0, 0                                           # 주사위 위치
dice = [1, 6, 3, 4, 5, 2]                                       # 초기 주사위 (위, 아래, 동, 서, 남, 북)
direction = 0                                                   # 초기 방향은 동쪽
result = 0

dx = [0, 1, 0, -1]         # 동남서북
dy = [1, 0, -1, 0]

def rotate_dice(dice, direction):
    if direction == 0:      # 동
        dice[2], dice[1], dice[3], dice[0] = dice[0], dice[2], dice[1], dice[3]
    elif direction == 1:    # 남
        dice[4], dice[1], dice[5], dice[0] = dice[0], dice[4], dice[1], dice[5]
    elif direction == 2:    # 서
        dice[3], dice[1], dice[2], dice[0] = dice[0], dice[3], dice[1], dice[2]
    elif direction == 3:    # 북
        dice[5], dice[1], dice[4], dice[0] = dice[0], dice[5], dice[1], dice[4]
    return dice

def bfs(B, start_x, start_y):
    visited = [[False for _ in range(M)] for _ in range(N)]
    count = 1
    queue = deque()
    queue.append((start_x, start_y))
    visited[start_x][start_y] = True

    while queue:
        (x, y) = queue.popleft()
        for i in range(4):
            nx, ny = x + dx[i], y + dy[i]
            if (0 <= nx < N) and (0 <= ny < M) and array[nx][ny] == B and not visited[nx][ny]:
                queue.append((nx, ny))
                visited[nx][ny] = True
                count += 1
    return count

for _ in range(K):
    nx, ny = dice_x + dx[direction], dice_y + dy[direction]

    # 이동방향에 칸이 없다면 반대방향으로 이동방향 설정
    if not((0 <= nx < N)and(0 <= ny < M)):
        direction += 2
        if direction == 4:
            direction = 0
        elif direction == 5:
            direction = 1
        # print(direction)
        nx, ny = dice_x + dx[direction], dice_y + dy[direction]

    # 주사위 굴리기
    dice = rotate_dice(dice, direction)
    A = dice[1]                             # 주사위 아랫면에 있는 정수
    B = array[nx][ny]                       # 주사위가 도착한 칸에 적힌 정수
    dice_x, dice_y = nx, ny

    # 이동방향 결정하기
    if A > B:                               # 90도 시계방향
        direction += 1
        if direction == 4:
            direction = 0
    elif A < B:                             # 90도 반시계방향
        direction -= 1
        if direction == -1:
            direction = 3

    # bfs로 연속한 칸 수 구하기
    C = bfs(B, nx, ny)
    result += B * C
    # print(B, C)

print(result)