코드트리 - 고대 문명 유적 탐사 (Python)

import sys
from collections import deque

# 제출 전 삭제
sys.stdin = open("input.txt")

# Solution
# Step 1: 유물 회전
def get_rotate_data(graph):
    result_list = []
    for i in range(3):
        for j in range(3):
            for d in range(3):
                new_graph = rotate(i, j, d, graph)
                point, new_graph = get_point(new_graph)
                result_list.append([point, d, i, j, new_graph])

    return result_list

# 회전
def rotate(x, y, d, graph):
    new_graph = [[0] * 5 for _ in range(5)]

    # 90도 회전
    if d == 0:
        for i in range(3):
            for j in range(3):
                new_graph[j + x][2 - i + y] = graph[x + i][y + j]
    # 180도 회전
    if d == 1:
        for i in range(3):
            for j in range(3):
                new_graph[2 - i + x][2 - j + y] = graph[x + i][y + j]
    # 270도 회전
    if d == 2:
        for i in range(3):
            for j in range(3):
                new_graph[2 - j + x][i + y] = graph[x + i][y + j]

    for i in range(5):
        for j in range(5):
            if new_graph[i][j] == 0:
                new_graph[i][j] = graph[i][j]

    return new_graph

# 점수 획득
def get_point(graph):
    visited = [[False] * 5 for _ in range(5)]
    total_point = 0

    for i in range(5):
        for j in range(5):
            if not visited[i][j]:
                point, graph = bfs(i, j, graph, visited)
                total_point += point

    return total_point, graph

# search
def bfs(r, c, graph, visited):
    target = graph[r][c]

    dq = deque()
    dq.append([r, c])  # x, y
    visited[r][c] = True

    nodes = []
    nodes.append([r, c])

    while dq:
        x, y = dq.popleft()

        for dx, dy in dirs:
            nx, ny = x + dx, y + dy

            if in_range(nx, ny) and not visited[nx][ny] and graph[nx][ny] == target:
                dq.append([nx, ny])
                nodes.append([nx, ny])
                visited[nx][ny] = True

    if len(nodes) >= 3:
        for x, y in nodes:
            graph[x][y] = 0
    point = len(nodes) if len(nodes) >= 3 else 0

    return point, graph

# Step 2: 격자 선택
def select_best_graph(result_list):
    # result_list : [point, degree, i, j, graph]
    result_list.sort(key=lambda x: (-x[0], x[1], x[3], x[2]))
    best_result = result_list[0]
    return best_result[0], best_result[4]

# Step 3: 유물 추가하기
def add_new_number(graph):
    global num_idx

    for col in range(5):
        for row in range(4, -1, -1):
            if graph[row][col] == 0:
                graph[row][col] = nums[num_idx]
                num_idx += 1

    return graph

# Step 4: 유물 연쇄 획득
def get_additional_points(graph):
    additional_point = 0

    point, graph = get_point(graph)
    if point > 0:
        graph = add_new_number(graph)
        additional_point += point

    while point > 0:
        point, graph = get_point(graph)
        if point > 0:
            graph = add_new_number(graph)
            additional_point += point

    return additional_point, graph

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

# Main
K, M = map(int, sys.stdin.readline().split())
maps = []
nums = []  # 벽면 유물
num_idx = 0  # 벽면 유물 index
point_list = []  # 획득한 포인트
dirs = [(1, 0), (-1, 0), (0, 1), (0, -1)]

# get maps
for _ in range(5):
    maps.append(list(map(int, sys.stdin.readline().split())))

# get nums:
nums = list(map(int, sys.stdin.readline().split()))

# game
for _ in range(K):
    total_point = 0
    rotate_result_list = get_rotate_data(maps)
    point, maps = select_best_graph(rotate_result_list)

    if point == 0:
        break
    total_point += point
    maps = add_new_number(maps)

    # 연쇄 반응
    point, maps = get_additional_points(maps)

    # 점수 계산
    total_point += point
    point_list.append(total_point)

print(*point_list)