[Python/Codetree] 마법사의 숲 탐색

https://www.codetree.ai/ko/frequent-problems/samsung-sw/problems/magical-forest-exploration/description

코딩테스트 기출 문제 설명: 마법의 숲 탐색 | 코드트리

문제

정령들이 $R$헹 $C$열의 격자 형태로 이루어진 마법의 숲을 탐색하려고 합니다. 격자는 가장 위를 1행, 가장 아래를 $R$행으로 합니다.

숲의 동쪽, 서쪽, 남쪽은 마법의 벽으로 막혀 있으며, 정령들은 숲의 북쪽을 통해서만 숲에 들어올 수 있습니다.

총 $K$명의 정령은 각자 골렘을 타고 숲을 탐색합니다. 각 골렘은 십자 모양의 구조를 가지고 있으며, 중앙 칸을 포함해 총 5칸을 차지합니다. 골렘의 중앙을 제외한 4칸 중 한 칸은 골렘의 출구입니다. 정령은 어떤 방향에서든 골렘에 탑승할 수 있지만 골렘에서 내렬 때에는 정해진 출구를 통해서만 내릴 수 있습니다.

$i$번째로 숲을 탐색하는 골렘은 숲의 북쪽 방향의 바깥에서 시작해 골렘이 중앙의 $c_{i}$열이 되도록 하는 위치에서 내려오기 시작합니다. 초기 골렘의 출구는 $d_{i}$(0 ~ 3으로 주어지며, 순서대로 북, 동, 남, 서에 대응)의 방향에 위치해 있습니다.

골렘은 숲을 탐색하기 위해 다음과 같은 우선순위로 이동합니다.(처음 시작 위치만 숲의 바깥에서도 동일하게 적용됩니다.) 더 이상 움직이지 못할 때까지 해당 과정을 반복합니다.

 

(1) 남쪽으로 한 칸 내려갑니다.

위 그림의 초록색 칸들이 비어있을 때에만 남쪽으로 내려갈 수 있습니다.

 

(2) (1)의 방법으로 이동할 수 없으면 서쪽 방향으로 회전하면서 내려갑니다.

위 그림의 초록색 칸들이 비어있을 때에만 서쪽 방향으로 회전하면서 내려갈 수 있습니다. 이 때 서쪽 방향으로 한 칸 이동을 한 뒤 내려가기 때문에 골렘을 기준으로 서쪽 한 칸이 모두 비어 있어야 함에 유의합니다. 또한 이렇게 이동할 때 출구가 반시계 방향으로 이동합니다.

 

(3) (1)과 (2)의 방법으로 이동할 수 없으면 동쪽 방향ㅇ으로 회전하면서 내려갑니다.

위 그림의 초록색 칸들이 비어있을 때에만 동쪽 방향으로 회전하면서 내려갈 수 있습니다. 이 때 동쪽 방향으로 한 칸 이동을 한 뒤 내려가기 때문에 골렘을 기준으로 동쪽 한 칸이 모두 비어 있어야 함에 유의합니다. 또한 이렇게 이동할 때 출구가 시계 방향으로 이동합니다.

골렘이 이동할 수 있는 가장 남쪽에 도달해 더 이상 이동할 수 없으면 정령은 골렘 내에서 상하좌우 인접한 칸으로 이동이 가능합니다. 단, 만약 현재 위치하고 있는 골렘의 출구가 다른 골렘과 인접하고 있다면 해당 출구를 통해 다른 골렘으로 이동할 수 있습니다.

정령은 갈 수 있는 모든 칸 중 가장 남쪽의 칸으로 이동하고 이동을 환전히 종료합니다. 이때 정령의 위치가 해당 정령의 최종 위치가 됩니다. 아래 그림에서는 정령이 출구를 타고 6번째 행까지 이동이 가능합니다.

이 문제에서는 정령의 최종 위치의 행 번호의 합을 구해야 하기에 각 정령이 도달하게 되는 최종 위치를 누적해야 합니다.

아래의 그림처럼 만약 골렘이 최대한 남쪽으로 이동했지만 골렘의 몸 일부가 여전히 숲을 벗어난 상태라면, 해당 골렘을 포함해 숲에 위치한 모든 골렘들은 숲을 빠져나간 뒤 다음 골렘부터 새롭게 숲의 탐색을 시작합니다. 단, 이 경우에는 정령이 도달하는 최종 위치를 답에 포함시키지 않습니다.

즉, 아래와 같이 숲이 텅 빈 상태에서 다음 골렘이 탐색을 시작합니다.

골렘들이 숲에 진입함에 따라 각 정령들이 최종적으로 위치한 행의 총합을 구하는 프로그램을 작성하세요. 숲이 다시 텅 비게 돼도 행의 총합은 누적되는 것에 유의합니다.

풀이

Code

import sys
from typing import List, Tuple, Optional
from collections import deque
input = sys.stdin.readline

# 1. 숲 클래스 정의
class Forest:
    # 1-1.
    def __init__(self):
        # 1-1-1. 숲 그래프 선언
        self.graph = self._create_graph()
        # 1-1-2. 총합 변수 선언
        self.summation = 0
        # 1-1-3. 골렘 이동 방향 리스트 선언
        self.dirs = [[(1, 0)], # 남
                     [(0, -1), (1, 0)], # 남서
                     [(0, 1), (1, 0)] # 남동
                    ]

    # 1-2. 숲 그래프 생성 함수 정의
    @staticmethod
    def _create_graph() -> List[List[int]]:
        # 1-2-1. 빈 숲 그래프 반환
        return [[0 for _ in range(C)] for _ in range(R)]

    # 1-3. 이동 가능 여부 체크 함수 정의
    def _check(self,
               golem_index:List[Tuple],
               dir:List[Tuple[int]]) -> Optional[List[Tuple]]:
        # 1-3-1.
        next_golem_index = golem_index[:]
        for dir_x, dir_y in dir:
            # 이동 후 위치 인덱스 생성
            next_golem_index = [(x+dir_x, y+dir_y) for x, y in next_golem_index]
            # 범위를 벗어나거나 골렘이 이미 위치한 인덱스인 경우 None 반환
            for x, y in next_golem_index:
                if y<0 or y>=C or x>=R or (0<=x<R and self.graph[x][y]) :
                    return None
        # 1-3-2. 이동이 가능한 경우 이동 후 위치 인덱스 반환
        return next_golem_index

    # 1-4. 골렘 이동 함수 정의
    def _move(self,
              start_col:int,
              exit_dir:int) -> Tuple[List[Tuple], int]:
        # 1-4-1. 서브 골렘 위치 인덱스 생성
        sub_golem_index = None
        # 1-4-2. 골렘 위치 인덱스 생성
        golem_index = [(-3, start_col), # 북
                       (-2, start_col+1), # 동
                       (-1, start_col), #남
                       (-2, start_col-1), #서
                       (-2, start_col) #중앙
                      ]
        # 1-4-3.
        while golem_index != sub_golem_index:
            # 서브 골렘 위치 인덱스 업데이트
            sub_golem_index = golem_index[:]
            for idx, dir in enumerate(self.dirs):
                # 이동 가능 여부가 None 값이 아닌 경우
                if (checked:=self._check(golem_index, dir)) is not None:
                    # 골렘 위치 인덱스 업데이트
                    golem_index = checked[:]
                    # 이동 방향에 따른 출구 방향 업데이트
                    # 남서쪽 회전의 경우
                    if idx == 1:
                        exit_dir = (exit_dir-1)%4
                    # 남동쪽 회전의 경우
                    elif idx == 2:
                        exit_dir = (exit_dir+1)%4

                    break
        # 1-4-4. 최종 위치 및 출구 방향 반환
        return golem_index[:], exit_dir

    # 1-5. 골렘 배치 함수 정의
    def _placement(self,
                   num:int,
                   start_col:int,
                   exit_dir:int) -> Optional[Tuple]:
        # 1-5-1. 골렘 이동 후 위치 인덱스 및 출구 방향 추출
        golem_index, exit_dir = self._move(start_col, exit_dir)
        # 1-5-2. 골렘이 숲에 완전히 들어오지 못한 경우
        if min([x for tup in golem_index for x in tup]) < 0:
            # 숲 그래프 초기화
            self.graph = self._create_graph()

            # None 반환
            return None
        # 1-5-3. 골렘이 숲에 완전히 들어온 경우
        else:
            # 숲 그래프 업데이트
            for idx, index in enumerate(golem_index):
                x, y = index
                if exit_dir == idx:
                    self.graph[x][y] = -num
                else:
                    self.graph[x][y] = num

            # 정령 위치 반환
            return golem_index[-1]

    # 1-6. 정령 이동 함수 정의
    def _spirit_movement(self,
                         num:int,
                         split_index:Tuple[int]) -> int:
        # 1-6-1. 큐 생성 후 (골렘번호, 정령 위치) 삽입
        x, y = split_index
        queue = deque([(num, x, y)])
        # 1-6-2. 방문 여부 리스트 생성 후 정령 위치 방문 처리
        visited = [[False for _ in range(C)] for _ in range(R)]
        visited[x][y] = True
        # 1-6-3. 최대 행 변수 초기화
        max_row = -1
        # 1-6-4.
        while queue:
            # 골렘 번호 및 정령 위치 추출
            num, x, y = queue.popleft()
            for dir_x, dir_y in [(-1, 0), (0, 1), (1, 0), (0, -1)]:
                # 다음 위치 설정
                nx, ny = x+dir_x, y+dir_y
                # 예외 처리
                if nx<0 or nx>=R or ny<0 or ny>=C: continue
                # 같은 골렘 번호이거나 현재 위치가 출구인 경우
                if num==abs(self.graph[nx][ny]) or (num<0 and self.graph[nx][ny]):
                    # 다음 위치에 방문하지 않은 경우
                    if not visited[nx][ny]:
                        # 방문 처리
                        visited[nx][ny] = True
                        # 골렘 최대 행 위치 업데이트
                        max_row = max(max_row, nx)
                        # 큐에 (골렘 번호, 다음 정령 위치) 삽입
                        queue.append((self.graph[nx][ny], nx, ny))
        # 1-6-5. 최대 행 변수 반환
        return max_row+1

    # 1-7. 실행 함수 정의
    def execute(self,
                num:int,
                start_col:int,
                exit_dir:int) -> None:
        # 1-7-1. 골렘 배치 후 정령 위치 추출
        split_index = self._placement(num, start_col, exit_dir)
        # 1-7-2. 정령 위치가 None이 아닌 경우
        if split_index is not None:
            # 정령 이동 후 총합 변수 업데이트
            self.summation += self._spirit_movement(num, split_index)

R, C, K = map(int, input().split())
# 2. 숲 Object 생성
forest = Forest()
# 3.
for num in range(1, K + 1):
    c, d = map(int, input().split())
    # 3-1. 실행
    forest.execute(num=num,
                   start_col=c-1,
                   exit_dir=d)
# 4. 결과 출력
print(forest.summation)