코드트리 - 꼬리잡기 (Python)

import sys
from collections import deque

# Solution
# Step 1: 1칸 전진
def move_one_step():
    visited = [[False] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            if graph[i][j] == 1 and not visited[i][j]:
                rotate([i, j], visited)

def rotate(head, visited):
    head, tail, _ = get_head_tail(*head)

    # move tail
    r, c = tail
    for dr, dc in dirs:
        nr, nc = r + dr, c + dc
        if in_range(nr, nc) and graph[nr][nc] == 2:
            graph[nr][nc] = 3
            graph[r][c] = 4
            break

    # move head
    r, c = head
    for dr, dc in dirs:
        nr, nc = r + dr, c + dc
        if in_range(nr, nc) and graph[nr][nc] == 4:
            graph[nr][nc] = 1
            graph[r][c] = 2
            visited[nr][nc] = True
            break

# Step 2: 공 던지기
def throw(round):
    start, dir = get_start_and_dir(round)
    x, y = start

    if graph[x][y] in [1, 2, 3]:
        return x, y

    nx, ny = x + dir[0], y + dir[1]
    while in_range(nx, ny) and not graph[nx][ny] in [1, 2, 3]:
        nx += dir[0]
        ny += dir[1]

    return nx, ny

def get_start_and_dir(round):
    round %= 4 * n
    if 0 <= round < n:
        start = [round, 0]
        dir = [0, 1]
    elif n <= round < 2 * n:
        start = [n - 1, round - n]
        dir = [-1, 0]
    elif 2 * n <= round < 3 * n:
        start = [3 * n - 1 - round, n - 1]
        dir = [0, -1]
    else:
        start = [0, 4 * n - 1 - round]
        dir = [1, 0]

    return start, dir

def get_head_tail(x, y):
    visited = [[False] * n for _ in range(n)]
    dq = deque()
    dq.append([x, y, 1])
    visited[x][y] = True

    if graph[x][y] == 1:
        head = [x, y]
        dist = 1
    elif graph[x][y] == 3:
        tail = [x, y]

    while dq:
        r, c, d = dq.popleft()

        for dr, dc in dirs:
            nr, nc = r + dr, c + dc
            if in_range(nr, nc) and not visited[nr][nc]:
                if graph[nr][nc] == 2:
                    dq.append([nr, nc, d + 1])
                    visited[nr][nc] = True
                elif graph[nr][nc] == 1:
                    head = [nr, nc]
                    dist = d + 1
                elif graph[nr][nc] == 3:
                    tail = [nr, nc]

    return head, tail, dist

def get_distance(head, target):
    dq = deque()
    dq.append(head + [1])

    while dq:
        r, c, d = dq.popleft()

        if r == target[0] and c == target[1]:
            return d

        for dr, dc in dirs:
            nr, nc = r + dr, c + dc
            if in_range(nr, nc):
                if d == 1 and graph[nr][nc] == 2:
                    dq.append([nr, nc, d + 1])
                if d != 1 and graph[nr][nc] in [2, 3]:
                    dq.append([nr, nc, d + 1])
    return 0

def swap(x1, y1, x2, y2):
    graph[x1][y1], graph[x2][y2] = graph[x2][y2], graph[x1][y1]

def in_range(x, y):
    return 0 <= x < n and 0 <= y < n

# Main
n, m, k = map(int, sys.stdin.readline().split())
graph = []
for _ in range(n):
    graph.append(list(map(int, sys.stdin.readline().split())))

dirs = [[1, 0], [-1, 0], [0, 1], [0, -1]]

score = 0
for round in range(k):
    move_one_step()
    x, y = throw(round)
    if in_range(x, y):
        head, tail, dist = get_head_tail(x, y)
        # dist = get_distance(head, [x, y])
        score += pow(dist, 2)
        swap(*head, *tail)
print(score)