백준 2302번 : 극장 좌석 in Python

2302번: 극장 좌석

 

www.acmicpc.net

이 문제는 한 줄로 되어있는 좌석에 사람들이 앉는 서로 다른 방법의 수를 구하는 문제이다. 이때 VIP 회원은 반드시 자기 좌석에만 앉아야 하고, 일반 회원은 양 옆의 사람끼리만 자리를 바꿀 수 있다. 예를 들어 5번 좌석에 앉은 일반 회원은 4번 좌석이나 6번 좌석에 앉을 수 있다. 하지만 3번 좌석이나 7번 좌석에는 앉을 수 없다. 

이를 해결하기 위해서는 dynamic programming을 이용한다. 우선 각 i번째 좌석마다 그 좌석을 예약한 사람이 그 자리에 그대로 앉는 경우와, 예약한 사람 오른쪽에 있는 사람과 바꿔 앉는 경우를 각각 저장하는 배열을 i번째 entry로 하는 배열 dp를 만든다. 초기값으로는 모두 0을 지정하되 dp[1][0]은 1로 하고, 1번째 좌석을 VIP 회원이 예약하지 않은 경우에 한해 dp[1][1]을 1로 한다. 이는 1번째 좌석이 VIP 좌석이 아닌 경우 2번째 좌석을 예약한 사람과 자리를 바꿀 수 있기 때문이다.

이제 2번째 좌석부터 N번째 좌석까지 순회한다. 만약 i번째 좌석을 VIP 회원이 예약한 경우에는 왼쪽 좌석, 오른쪽 좌석을 예약한 사람과 자리를 바꿀 수 없으므로 dp[i][0]은 dp[i-1][0]과 같고, dp[i][1]은 0이다. 반대로 i번째 좌석을 일반 회원이 예약한 경우 양 옆 좌석과 자리를 바꿀 수 있으므로 dp[i][0]은 dp[i-1][0]과 dp[i-1][1]의 합과 같고, dp[i][1]은 dp[i-1][0]과 같다. 이때 dp[i][1]이 dp[i-1][0]만 반영하는 이유는 만약 왼쪽 좌석과 자리를 바꾼 경우에는 오른쪽 좌석과 자리를 바꾸는 것이 불가능하기 때문이다. 이를 i = N일 때까지 계산하면 dp[N][0]에는 모든 좌석에 사람들이 앉는 서로 다른 방법의 수가 저장된다.

이를 코드로 구현하면 다음과 같다.

N = int(input())
# VIP 회원인지를 판별하기 위한 배열
VIP = [False for _ in range(N+1)]
M = int(input())
for _ in range(M):
    n = int(input())
    VIP[n] = True
# 오른쪽 좌석과 자리를 바꾸지 않는 경우와 바꾸는 경우의 수를 각 좌석마다 저장하는 배열
dp = [[0, 0] for _ in range(N+1)]
# 1번째 좌석을 VIP 회원이 예약하지 않은 경우에만 dp[1][1]을 1로 한다.
if not VIP[1]:
    dp[1][1] = 1
dp[1][0] = 1
for i in range(2, N+1):
    # i번째 좌석을 VIP 회원이 예약한 경우
    if VIP[i]:
        dp[i][0] = dp[i-1][0]
    else:
        dp[i][0] = dp[i-1][0] + dp[i-1][1]
        dp[i][1] = dp[i-1][0]
print(dp[N][0])