[python] 코드트리 - 꼬리잡기놀이

문제

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

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

 

n * n 격자에서 꼬리잡기 놀이를 합니다. 입력형식은 아래 그림을 참고바랍니다. 

각 팀은 머리사람(맨 앞 사람)과 꼬리사람(맨 뒤 사람)으로 구성되어 있습니다.

• 각 팀은 게임에서 주어진 이동 경로만을 따라 이동하며, 이 경로는 팀별로 겹치지 않도록 보장됩니다.
• 라운드가 진행되면 각 팀은 머리사람을 따라 한 칸씩 이동합니다.
• 라운드마다 공이 특정 경로를 따라 던져지며, 공이 이동하는 경로에서 사람이 있을 경우, 공은 최초로 만나는 사람에게 전달됩니다. 공을 받은 사람은 점수를 얻게 됩니다.
• 점수는 그 사람이 팀 내에서 머리사람부터 k번째 사람일 때, k의 제곱(k²)만큼 주어집니다.
• 공을 획득한 팀은 머리사람과 꼬리사람의 위치를 바꾸어 이동 방향이 반대로 변경됩니다.

공이 던져지는 순서는 다음과 같습니다.

 

 

 

풀이

 

1. 배열 처리

배열에서 빈칸은 0, 경로는 4, 그리고 사람이 있는 곳은 1, 2, 3으로 표시되며, 공이 날아가는 경로에서 사람을 만나면 점수를 획득할 수 있습니다. 중요한 점은, 사람의 위치가 머리, 중간, 꼬리와 상관없이 공이 맞으면 해당 팀원이 점수를 얻는다는 것입니다.

따라서 사람의 위치를 세부적으로 구분할 필요가 없다는 점을 고려해, 배열 내 팀원들을 5 이상의 숫자로 표시하여 팀 소속만 구분했습니다. 이렇게 함으로써 배열 내에서 팀 경로와 팀원을 효율적으로 관리할 수 있었습니다

 

 

2. 팀 관리 및 경로 탐색 (BFS 사용)

 

각 팀을 관리하기 위해 딕셔너리 형태로 팀 번호: [팀의 경로와 팀원 위치]를 저장했습니다. 팀 경로와 팀원의 위치를 알아내기 위해 BFS(너비 우선 탐색)를 활용하여 경로를 탐색했습니다.

BFS는 처음 머리사람의 위치에서 시작해 경로를 따라가면서 팀원들의 좌표를 모두 탐색합니다. 이렇게 각 팀의 팀원 위치를 효율적으로 찾고, 이를 딕셔너리에 저장해 관리합니다.

def get_group(arr, i, j):
    # BFS를 활용해 팀원 위치 탐색
    ...​

 

 

3. 공 던지기 및 점수 계산

 

공이 던져지면, 코드에서는 shoot_ball 함수를 통해 공이 맞는 경로에 사람이 있는지 확인합니다. 사람이 있다면, 해당 팀원의 위치를 기반으로 점수를 계산합니다.

이때 점수는 해당 팀원의 위치에 따라 달라지며, 팀 내에서 k번째 사람일 경우 점수는 k^2으로 계산됩니다. 공을 획득한 팀은 이후 머리사람과 꼬리사람의 위치를 바꾸어 이동 방향이 변경됩니다.

def shoot_ball(arr, groups, turn):
    # 공이 던져지는 경로에서 사람을 만나면 점수를 계산
    ...

 

 

4. 팀 이동 처리

 

move_group 함수는 각 라운드마다 팀의 이동을 처리하는 역할을 합니다. 팀의 이동은 머리사람 부터가아닌 꼬리 사람 부터 이동합니다. 그 이유는 만약 팀의 머리사람과 꼬리사람이 이어저 있다면, 머리 사람은 이동 할 칸을 찾을 수 없기 때문입니다.

공을 획득한 팀은 머리사람과 꼬리사람의 위치가 뒤바뀌고, 이동 방향도 반대가 됩니다.  이를 통해 공이 맞은 팀은 이후 라운드에서 방향이 전환됩니다.

def move_group(arr, groups):
    # 각 팀의 머리사람 이동, 꼬리사람은 그 자리를 따라감
    ...

 

 

 

 

 

전체 코드

import sys
from collections import deque, defaultdict
from typing import List, Tuple

# 입력을 받는 함수
def read_input() -> Tuple[int, int, int, List[List[int]]]:
    n, m, k = list(map(int, sys.stdin.readline().split()))  # n, m, k 값 입력받기
    arr = [list(map(int, input().split())) for _ in range(n)]  # 2D 배열 입력받기
    return n, m, k, arr


# 특정 좌표에 대한 점수 계산 함수
def calculate_score(group, y, x):
    score = 0
    # 해당 좌표의 점수를 계산
    for idx, (yy, xx) in enumerate(group):
        if (yy, xx) == (y, x):
            score = (idx + 1) ** 2  # 점수는 인덱스의 제곱
            break
    group = group[::-1]  # 그룹을 회전
    return score, group

# 공을 던져 맞는 그룹을 찾고 점수를 계산하는 함수
def shoot_ball(arr, groups, turn):
    n = len(arr)
    line = turn % n
    direction = (turn % (4 * n)) // n  # 방향을 결정하는 변수

    def hit_check(i, j, is_vertical):
        # 공이 맞는지 확인하는 내부 함수
        if arr[i][j] > 4:
            group_num = arr[i][j]
            score, group = calculate_score(groups[group_num], i, j)
            groups[group_num] = group  # 그룹을 업데이트
            return score, groups
        return 0, groups

    if direction == 0:  # 우
        for j in range(n):
            score, groups = hit_check(line, j, False)
            if score > 0:
                return score, groups
    elif direction == 1:  # 상
        for i in range(n - 1, -1, -1):
            score, groups = hit_check(i, line, True)
            if score > 0:
                return score, groups
    elif direction == 2:  # 좌
        line = n - line - 1
        for j in range(n - 1, -1, -1):
            score, groups = hit_check(line, j, False)
            if score > 0:
                return score, groups
    elif direction == 3:  # 하
        line = n - line - 1
        for i in range(n):
            score, groups = hit_check(i, line, True)
            if score > 0:
                return score, groups

    return 0, groups  # 맞는 그룹이 없으면 0점 반환

# 그룹을 이동시키는 함수
def move_group(arr, groups):
    dx = [0, 1, -1, 0]  # x축 이동 방향
    dy = [1, 0, 0, -1]  # y축 이동 방향
    n = len(arr)

    for team_num in groups.keys():
        team = groups[team_num]

        # 꼬리 좌표를 미리 4로 변경
        ty, tx = team[-1]
        arr[ty][tx] = 4

        # 머리가 이동할 좌표 구하기
        y, x = team[0]
        for d in range(4):
            hy = y + dy[d]
            hx = x + dx[d]
            if 0 <= hy < n and 0 <= hx < n and arr[hy][hx] == 4:
                arr[hy][hx] = team_num  # 새로운 머리 좌표 설정
                break

        # 팀 좌표 업데이트 (머리 이동, 꼬리 제거)
        groups[team_num] = [(hy, hx)] + team[:-1]

    return arr, groups

# 특정 좌표에서 그룹을 추출하는 함수
def get_group(arr, i, j):
    n = len(arr)
    visited = [[0] * n for _ in range(n)]  # 방문 체크 배열

    dx = [0, 1, -1, 0]  # x축 이동 방향
    dy = [1, 0, 0, -1]  # y축 이동 방향

    group = [(i, j)]
    queue = deque([(i, j)])  # BFS 탐색을 위한 큐
    visited[i][j] = 1

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

        for d in range(4):
            ny, nx = y + dy[d], x + dx[d]

            if 0 <= ny < n and 0 <= nx < n and not visited[ny][nx]:
                if arr[ny][nx] == 2 or (arr[y][x] == 2 and arr[ny][nx] == 3):
                    group.append((ny, nx))
                    queue.append((ny, nx))
                    visited[ny][nx] = 1

    return group

# 메인 실행 함수
def main():
    n, m, k, arr = read_input()  # 입력 데이터 받기

    answer = 0
    team_num = 5  # 그룹 번호는 5번부터 시작
    groups = defaultdict(list)

    # 그룹 탐색 및 팀 번호 부여
    for i in range(n):
        for j in range(n):
            if arr[i][j] == 1:
                group = get_group(arr, i, j)
                groups[team_num] = group
                team_num += 1

    # 기존 배열에 팀 번호 표시
    for team_num in groups.keys():
        for y, x in groups[team_num]:
            arr[y][x] = team_num

    # 턴마다 게임 진행
    for turn in range(k):
        arr, groups = move_group(arr, groups)  # 그룹 이동
        score, groups = shoot_ball(arr, groups, turn)  # 공 던지기 및 점수 계산
        answer += score  # 점수 누적

    print(answer)  # 최종 점수 출력

# 프로그램 실행
if __name__ == "__main__":
    main()