코드트리 꼬리잡기놀이 (파이썬)

https://www.codetree.ai/training-field/frequent-problems/problems/tail-catch-play

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

머리와 꼬리가 붙어있어서 경로에 4가 없는 케이스도 생각해야한다

import sys
from collections import deque

input = sys.stdin.readline

# 공 던지는 방향 순서 고려해서 우상좌하 순서
dx = [0, -1, 0, 1]
dy = [1, 0, -1, 0]

n, m, k = map(int, input().split())

a = [list(map(int, input().split())) for _ in range(n)]

# 머리부터 시작해서 이어져있는 사람들을 모아 리스트 team을 만들고 teams에 넣는다
teams = []
for i in range(n):
    for j in range(n):
        if a[i][j] == 1:
            check = [[0 for _ in range(n)] for _ in range(n)]
            check[i][j] = 1
            team = [[i, j]]
            head, tail = 1, 0  # 현재 좌표가 머리 or 꼬리
            x, y = i, j
            while True:
                for d in range(4):
                    nx = x + dx[d]
                    ny = y + dy[d]
                    if 0 <= nx < n and 0 <= ny < n:
                        if a[nx][ny] == 3:
                            if head == 1:  # 머리에서 꼬리로 이동 방지
                                continue
                            else:
                                team.append([nx, ny])
                            tail = 1
                            break
                        elif a[nx][ny] == 2 and check[nx][ny] == 0:
                            team.append([nx, ny])
                            break
                if tail == 1:
                    break
                x, y = nx, ny
                check[x][y] = 1
                head = 0
            teams.append(team)

# 머리와 꼬리가 이어져있는지 검사
# circle 값이 1이면 해당 인덱스의 팀은 머리와 꼬리가 이어짐
circle = [0 for _ in range(m)]
for i, team in enumerate(teams):
    hx, hy = team[0]
    tx, ty = team[-1]
    if abs(hx - tx) + abs(hy - ty) == 1:
        circle[i] = 1

ball_dir = 0
ans = 0
for r in range(k):
    # 한 칸씩 이동
    for i, team in enumerate(teams):
        hx, hy = team[0]
        temp_team = []
        # 머리가 이동할 다음칸을 찾는다
        # 꼬리와 이어졌으면 다음칸이 현재 꼬리 칸
        for d in range(4):
            nx = hx + dx[d]
            ny = hy + dy[d]
            if 0 <= nx < n and 0 <= ny < n:
                if circle[i] == 1 and a[nx][ny] == 3:
                    temp_team.append([nx, ny])
                    break
                elif circle[i] == 0 and a[nx][ny] == 4:
                    temp_team.append([nx, ny])
                    break
        # 나머지 사람들의 좌표를 temp_team에 넣는다
        temp_team.extend(team[:-1])

        # 격자의 값을 수정해준다
        for j, (x, y) in enumerate(temp_team):
            if j == 0:
                a[x][y] = 1
            elif j == len(temp_team) - 1:
                a[x][y] = 3
            else:
                a[x][y] = 2
        # 머리와 꼬리가 안이어졌으면 마지막 칸을 경로가 있는 빈칸으로 수정
        if circle[i] == 0:
            x, y = team[-1]
            a[x][y] = 4

        # 한 칸씩 이동한 팀으로 업데이트
        teams[i] = temp_team

    # 공 던지기
    # ball_dir: 공을 던지는 방향
    # ball_start: 공이 시작되는 좌표
    if ball_dir == 0 or ball_dir == 1:
        ball_start = r % n
    else:
        ball_start = (n - 1) - (r % n)

    # bx, by: 공 시작 좌표
    if ball_dir == 0:
        bx, by = ball_start, -1
    elif ball_dir == 1:
        bx, by = n, ball_start
    elif ball_dir == 2:
        bx, by = ball_start, n
    else:
        bx, by = -1, ball_start

    # 공에 맞으면 점수 갱신하고 해당 팀의 머리와 꼬리 서로 교체
    while True:
        bx += dx[ball_dir]
        by += dy[ball_dir]
        if not 0 <= bx < n or not 0 <= by < n:
            break
        if 0 < a[bx][by] < 4:
            for i, team in enumerate(teams):
                for j, (x, y) in enumerate(team):
                    if [x, y] == [bx, by]:
                        hx, hy = team[0]
                        tx, ty = team[-1]
                        a[hx][hy], a[tx][ty] = a[tx][ty], a[hx][hy]
                        teams[i] = team[-1::-1]
                        ans += (j + 1) ** 2
            break

    # 공 던지는 방향 수정
    if r > 0:
        if ball_dir == 0 and ball_start == n - 1:
            ball_dir = (ball_dir + 1) % 4
        elif ball_dir == 1 and ball_start == n - 1:
            ball_dir = (ball_dir + 1) % 4
        elif ball_dir == 2 and ball_start == 0:
            ball_dir = (ball_dir + 1) % 4
        elif ball_dir == 3 and ball_start == 0:
            ball_dir = (ball_dir + 1) % 4

print(ans)