각 트랙의 움직임은 덱으로 구현하면 쉽다. 빈 트랙과 팀이 존재하는 트랙 두 덱을 관리하며, 이동 시 서로의 원소를 하나씩 구현하면 되기 때문. n의 크기가 충분히 작으므로, 덱들을 전부 참조하여 맵을 업데이트해도 크게 지장이 없다.
문제는 팀이 존재하는 트랙을 찾을 때이다. bfs로 탐색하든 dfs로 탐색하든, 트랙 전체를 팀이 메웠을 때는 예외 처리를 해주어야 한다(머리와 꼬리가 붙어있는 경우, 처음에 꼬리를 탐색하는 불상사를 일으킬 수 있다. 이 경우 잘못된 움직임으로 에러를 일으킨다). 다행히 머리 < 중간 < 꼬리 순으로 숫자가 정해져 있으므로, (탐색하고자 하는 위치의 숫자 - 원래 숫자) 가 1 이하의 차이를 만드는 지점만 탐색할 경우 머리부터 꼬리까지 차례대로 탐색할 수 있겠다.
덱의 또다른 장점은, 인덱스를 통해 어떤 지점에서 머리 지점까지의 거리를 쉽게 구할 수 있다는 점이다. 따라서 공을 던졌을 때 바로 스코어 연산이 가능하다.
공이 던져지는 좌표는 4n을 주기로 우상좌하 순으로 규칙적으로 생성되므로, 적절하게 구현해볼 수 있겠다.
풀이 코드
from collections import deque
import sys
input = sys.stdin.readline
dr = [0, -1, 0, 1]
dc = [1, 0, -1, 0]
n, m, k = map(int, input().split())
maps = [list(map(int, input().split())) for _ in range(n)]
visited = [[0]*n for _ in range(n)]
teams = [deque() for _ in range(m)]
left_tracks = [deque() for _ in range(m)]
scores = 0
def team_bfs(r, c, idx) :
visited[r][c] = idx+1
q = deque([(r, c)])
teams[idx].append((r, c))
left = False
while q :
r, c = q.popleft()
for i in range(4) :
ar, ac = r + dr[i], c + dc[i]
if -1 < ar < n and -1 < ac < n :
if not visited[ar][ac] and maps[ar][ac] not in [0, 4] and maps[ar][ac] - maps[r][c] <= 1 :
visited[ar][ac] = idx+1
q.append((ar, ac))
teams[idx].append((ar, ac))
if maps[r][c] == 3 and maps[ar][ac] == 4 :
lr, lc, left = ar, ac, True
if left :
left_bfs(lr, lc, idx)
def left_bfs(r, c, idx) :
visited[r][c] = idx+1
q = deque([(r, c)])
left_tracks[idx].append((r, c))
while q :
r, c = q.popleft()
for i in range(4) :
ar, ac = r + dr[i], c + dc[i]
if -1 < ar < n and -1 < ac < n :
if not visited[ar][ac] and maps[ar][ac] == 4 :
visited[ar][ac] = idx+1
q.append((ar, ac))
left_tracks[idx].append((ar, ac))
def init() :
cnt = 0
for i in range(n) :
for j in range(n) :
if maps[i][j] == 1 :
team_bfs(i, j, cnt)
cnt += 1
def move_all() :
for i in range(m) :
left_tracks[i].appendleft(teams[i].pop())
teams[i].appendleft(left_tracks[i].pop())
(hr, hc), (tr, tc) = teams[i][0], teams[i][-1]
maps[hr][hc], maps[tr][tc] = 1, 3
for j in range(1, len(teams[i])-1) :
r, c = teams[i][j]
maps[r][c] = 2
for r, c in left_tracks[i] :
maps[r][c] = 4
def throw_ball(i) :
global scores
i %= 4*n
s, d = i % n, i // n
if d == 0 :
r, c = s, 0
elif d == 1 :
r, c = n-1, s
elif d == 2 :
r, c = n-1-s, n-1
else :
r, c = 0, n-1-s
hit = False
while -1 < r < n and -1 < c < n :
if 1 <= maps[r][c] <= 3 :
hit = True
break
r += dr[d]
c += dc[d]
if not hit :
return
idx = visited[r][c]-1
scores += (teams[idx].index((r, c)) + 1) ** 2
teams[idx].reverse()
(hr, hc), (tr, tc) = teams[idx][0], teams[idx][-1]
maps[hr][hc], maps[tr][tc] = 1, 3
left_tracks[idx].reverse()
init()
for i in range(k) :
move_all()
throw_ball(i)
print(scores)