SWEA 2383ランチタイム(Python付き)

import sys
sys.stdin = open(r'/Users/kangdaeyoung/Downloads/sample_input (2).txt', 'r')

T = int(input())

for tc in range(1, T+1):
    N = int(input())
    brd = [list(map(int, input().split())) for _ in range(N)]

    answer = 1e10

    # 계단 위치 찾기 & 각 사람의 위치 찾기
    stairs = []
    people = []
    num_p = 0
    for r in range(N):
        for c in range(N):
            if brd[r][c] > 1:
                stairs.append([r, c, brd[r][c]])
            if brd[r][c] == 1:
                num_p += 1
                people.append([r, c])

    # 완전 탐색
    for i in range(1 << num_p):
        # 1. 사람 - 계단 매칭
        which_stair = [0] * num_p
        for j in range(num_p):
            if i & (1 << j):
                which_stair[j] = 1
        # 2. 각 사람이 계단을 배정 받았으면 그 계단과의 거리 계산
        dist_p_s = [0] * num_p
        for idx in range(num_p):
            # 마지막 원소 _ 는 걸리는 시간 의미
            stair_r, stair_c, _ = stairs[which_stair[idx]]
            person_r, person_c = people[idx]
            dist = abs(stair_r-person_r) + abs(stair_c-person_c)
            dist_p_s[idx] = dist
        # 3. 모든 사람이 계단을 통해서 빠져나올 때까지 시뮬레이션 go
        # 계단에 누구 있는지 정보
        stairs_info = [[-1, -1, -1],
                       [-1, -1, -1]]
        # 계단 입장을 위해 기다리는 친구 인원수 정보
        stairs_waiting = [0,
                          0]

        taken_time = 0
        while True:
            # 종료 조건
            # 글로벌 최소값보다 크다면 더 이상 찾아볼 가치 x
            if taken_time >= answer:
                break
            # 계단 - 사람 거리 원소가 모두 -1,
            # & waiting 비어있고,
            # & 계단 정보가 모두 -1 이면 "종료!"
            terminate_flag = True
            for idx in range(num_p):
                if dist_p_s[idx] >= -1:
                    terminate_flag = False

            for stair in range(2):
                if stairs_waiting[stair] > 0:
                    terminate_flag = False

            for stair in range(2):
                for position in range(3):
                    if stairs_info[stair][position] > 0:
                        terminate_flag = False

            if terminate_flag:
                break

            # 3.1) 현재 계단을 모두 내려온 사람 out
            for stair in range(2):
                for position in range(3):
                    if stairs_info[stair][position] == 0:
                        stairs_info[stair][position] -= 1

            # 3.2) 계단과의 거리 -1 인 애들 waiting line 으로 이동
            # 계단 도착 후, 1분 뒤에 내려갈 수 있단댜... ㅠ
            for idx in range(num_p):
                if dist_p_s[idx] == -1:
                    stair_idx = which_stair[idx]
                    stairs_waiting[stair_idx] += 1

            # 3.3) waiting 중인 애들 계단 입장
            for stair in range(2):
                flag = True
                while stairs_waiting[stair] and flag:
                    for idx in range(3):
                        if stairs_info[stair][idx] == -1:
                            stairs_info[stair][idx] = stairs[stair][2]
                            stairs_waiting[stair] -= 1
                            break
                    else:
                        flag = False

            # 3.4) 계단과의 거리 -1 & 계단에서 내려간 시간 -1
            for idx in range(num_p):
                if dist_p_s[idx] >= -1:
                    dist_p_s[idx] -= 1
            for stair in range(2):
                for position in range(3):
                    if stairs_info[stair][position] > 0:
                        stairs_info[stair][position] -= 1

            taken_time += 1

        if taken_time < answer:
            answer = taken_time

    print(f'#{tc} {answer}')