[Python 3] BOJ - 2305 "자리 배치"

 # 문제 링크

https://www.acmicpc.net/problem/2305

2305번: 자리 배치

 

 # 풀이

 개인적으로 이 문제의 핵심은 다이나믹 프로그래밍이라고 생각한다.

<그림 1. 자유석이 없는 경우>

 위 그림 1은 N이 4이고 자유석이 없는 경우를 나타낸 그림이다. 위 그림에서 물음표는 아직 주인이 정해지지 않은 자리이다. 따라서, 현재 그림 1은 3번 사람이 앉을 때 가능한 상황을 나타낸다고 할 수 있다. 왼쪽 위 그림의 경우는 1번 사람은 제자리에, 2번 사람 역시 제자리에 앉아 있는 경우이다. 이 때 3번 사람은 자신의 제자리인 3번 자리에 앉을 수도 있으나, 3번 자리의 오른쪽인 4번 자리에도 앉을 수 있다.

 왼쪽 아래 그림의 경우는 1번 사람은 제자리에, 2번 사람은 자신의 자리에서 오른쪽 자리에 앉아 있는 경우이다. 따라서 3번 사람은 제자리인 3번 자리를 제외한 왼쪽 좌석이나 오른쪽 좌석, 즉 2번 자리 아니면 4번 자리에 앉을 수 있다.

 오른쪽 위 그림의 경우는 1번 사람이 오른쪽, 2번 사람이 왼쪽 좌석에 앉은 경우이다. 따라서 3번 사람은 현재 물음표 자리인 3번 혹은 4번에 앉을 수 있으며, 바꿔 말해 제자리 혹은 오른쪽 좌석에 앉을 수 있다.

 마지막 경우인 오른쪽 아래 그림의 경우는 1번 사람과 2번 사람 모두 오른쪽에 앉은 경우이다. 이 때 3번 사람은 4번 자리에만 앉을 수 있다. 왜냐하면 현재 그림 1에서는 자유석이 없기 때문에 2칸 이상 차이나는 자리에는 앉을 수 없기 때문이다. 

<그림 2. 각 자리의 관계>

 위 그림 2는 현재 c번 사람이 앉을 때, 영향을 주는 사람을 나타내는 그림이다. c가 앉을 수 있는 좌석은 초록색 좌석이다. 그런데 좌석에 두 사람이 함께 앉아 있을 수 없다. 따라서, 현재 c가 왼쪽 좌석에 앉는다고 할 때 b가 제자리에 앉아 있거나, a가 오른쪽 빨간색에 앉아 있다면 c는 왼쪽 초록색 좌석에 앉을 수 없다. 만일 c가 제자리에 앉는다고 할 때도 역시 b가 오른쪽 파란 좌석에 앉아 있다면 불가능하다. 그러나, c가 오른쪽 좌석에 앉는 경우는 a와 b 모두 영향을 줄 수 없다. 따라서, 현재 c번 사람이 앉을 때 영향을 주는 사람은 이전에 앉았던 사람과 2번째 전에 앉았던 사람이다.

 따라서, 상태 공간을 현재 idx번째 사람이 앉을 차례이고, 2번째 전 사람이 앉았던 좌석을 ppast, 1번째 전 사람이 앉았던 좌석을 past, 자유석에 누군가가 앉아 있는지 아닌지 나타내는 flag일 때 가능한 경우의 수인 dp[idx][ppast][past][flag]인 4차원 상태공간으로 정의한다.

<그림 3. 3번 자리가 자유석인 경우>

 위 그림 3은 그림 1의 일부 중 자유석을 고려한 그림이다. 자유석에 해당하는 사람 번호는 문제에서 고려하지 않기 때문에 위 그림 3은 사실상 4번 사람이 자리를 고르는 과정과 같다. 따라서, 4번 사람은 제자리 혹은 왼쪽 좌석 중 하나를 선택해야 한다. 이 때, 왼쪽 좌석이 자유석이라고 가정하자. 그렇게 된다면 이 때는 왼쪽 좌석을 선택한 게 아닌 자유석을 선택한 것과 동일하다.

 만일 자유석을 선택한 것이 아니라 4번 사람이 왼쪽 좌석을 선택한 것이 올바르다고 가정해보자. 그렇게 되면 다음 사람은 자유석을 선택할 수 있다. 그런데, 이미 4번 사람이 왼쪽 좌석이자 동시에 자유석인 3번 자리에 앉아 있기 때문에 자유석에 두 사람이 앉을 수 있는 모순이 발생하게 된다. 따라서, 귀류법에 의해 모순이 발생하기 때문에 왼쪽 혹은 오른쪽 좌석이 자유석이라면 반드시 자유석에 앉은 것이라고 생각해야 올바른 정답을 구할 수 있다.

 따라서, 점화식은 아래와 같다.$$\begin{cases}
 & \text{ if flag is zero,} \ dp[idx][ppast][past][flag] += dp[idx + 1][past][3][1 - flag] \\ & \text{ if left seat is available,} \ dp[idx][ppast][past][flag] += dp[idx + 1][past][1][flag] \\ & \text{ if right seat is available,} \ dp[idx][ppast][past][flag] += dp[idx + 1][past][2][flag] \\ & \text{ if now seat is available,} \ dp[idx][ppast][past][flag] += dp[idx + 1][past][0][flag] \\ \end{cases}$$

 

 # 코드

import sys

# go : 현재 idx번째 사람이고, 2번째 이전 사람이 선택한 좌석이 ppast, 1번째 이전 사람이 선택한 
#      좌석이 past, 현재 자유석에 누가 앉아있는지 여부가 flag일 때 가능한 자리 배치의 경우의 수
#      를 리턴하는 함수
def go(idx, ppast, past, flag):
    # Base Case : 모두 앉았다면 가능한 경우이므로 1 리턴
    if idx == n - 1:
        return 1
    # Memoization
    if dp[idx][ppast][past][flag] != -1:
        return dp[idx][ppast][past][flag]
    dp[idx][ppast][past][flag] = 0
    # 자유석이 비어있는 경우
    if not flag:
        dp[idx][ppast][past][flag] += go(idx + 1, past, 3, 1 - flag)
    # 제자리에 앉을 수 있는 경우
    if past != 2:
        dp[idx][ppast][past][flag] += go(idx + 1, past, 0, flag)
    # 왼쪽 좌석이 자유석이 아니고 앉을 수 있는 경우
    if past != 0 and ppast != 2:
        if arr[idx] - 1 >= 1 and not check[arr[idx] - 1]:
            dp[idx][ppast][past][flag] += go(idx + 1, past, 1, flag)
    # 오른쪽 좌석이 자유석이 아니고 앉을 수 있는 경우
    if arr[idx] + 1 <= n and not check[arr[idx] + 1]:
        dp[idx][ppast][past][flag] += go(idx + 1, past, 2, flag)
    return dp[idx][ppast][past][flag]

# 입력부
n = int(sys.stdin.readline())
k = int(sys.stdin.readline())
arr = []
# check : 현재 좌석이 자유석인지 아닌지 체크하는 배열
check = [False] * (n + 1)
for i in range(1, n + 1):
    if i == k:
        check[i] = True
        continue
    arr.append(i)
    
# dp : 현재 idx번째 사람이고, 2번째 이전 사람이 선택한 좌석이 ppast, 1번째 이전 사람이 선택한 
#      좌석이 past, 현재 자유석에 누가 앉아있는지 여부가 flag일 때 가능한 자리 배치의 경우의 수
#      를 저장하는 4차원 상태 공간 배열
dp = [[[[-1] * 2 for _ in range(5)] for _ in range(5)] for _ in range(n + 1)]

# 정답 출력
print(go(0, -1, -1, 0))