[백준] 1309: 동물원 (Python)

1309: 동물원

문제

어떤 동물원에 가로로 두칸 세로로 N칸인 아래와 같은 우리가 있다.

이 동물원에는 사자들이 살고 있는데 사자들을 우리에 가둘 때, 가로로도 세로로도 붙어 있게 배치할 수는 없다. 이 동물원 조련사는 사자들의 배치 문제 때문에 골머리를 앓고 있다.

동물원 조련사의 머리가 아프지 않도록 우리가 2*N 배열에 사자를 배치하는 경우의 수가 몇 가지인지를 알아내는 프로그램을 작성해 주도록 하자. 사자를 한 마리도 배치하지 않는 경우도 하나의 경우의 수로 친다고 가정한다.

입력

첫째 줄에 우리의 크기 N(1≤N≤100,000)이 주어진다.

출력

첫째 줄에 사자를 배치하는 경우의 수를 9901로 나눈 나머지를 출력하여라.

풀이

직전 우리에 사자를 어떻게 배치했느냐에 따라 이번 우리에 사자를 어떻게 배치할 수 있는지가 달라진다. 따라서 dp테이블을 직전에 아무 사자도 배치하지 않은 경우, 왼쪽에 사자를 배치한 경우, 오른쪽에 사자를 배치한 경우 3가지 경우를 고려한 (n+1)*3 리스트로 생성했다. 이를 이용한 점화식은 다음과 같다.

1) dp[i][0] = sum(dp[i-1])

i번째에 아무 사자도 배치하지 않는 경우의 수 = i-1번째에 아무 사자도 배치하지 않은 경우 + i-1번째에 왼쪽에 사자를 배치한 경우+i-1번째에 오른쪽에 사자를 배치한 경우

2) dp[i][1] = (dp[i-1][0] + dp[i-1][2]) i번째에 왼쪽 우리에 사자를 배치하는 경우의 수 = i-1번째에 아무 사자도 배치하지 않은 경우 + i-1번째에 오른쪽에 사자를 배치한 경우

3) dp[i][2] = (dp[i-1][0] + dp[i-1][1]) 

i번째에 오른쪽 우리에 사자를 배치하는 경우의 수 = i-1번째에 아무 사자도 배치하지 않은 경우 + i-1번째에 왼쪽에 사자를 배치한 경우

 

메모리 초과를 막기 위해 9901로 나눈 결과를 출력하라고 되어있는데, 맨 처음 제출했던 코드는 for문의 점화식에만 mod로 나눠주고 최종 결과 출력 시 mod로 나눠주지 않았기 때문에 "틀렸습니다" 처리가 됐다. 반드시 for문과 print문 둘 다 mod로 나눠주도록 하자.

# 1309: 401 - 다이나믹 프로그래밍 1 연습 - 동물원
n = int(input())
mod = 9901
dp = [[0] * 3 for i in range(n+1)]
dp[1] = [1, 1, 1] # 한 마리도 배치하지 않은 경우, 왼쪽에 한 마리 배치한 경우, 오른쪽에 한 마리 배치한 경우
for i in range(2, n+1):
    dp[i][0] = sum(dp[i-1]) % mod
    dp[i][1] = (dp[i-1][0] + dp[i-1][2]) % mod
    dp[i][2] = (dp[i-1][0] + dp[i-1][1]) % mod
print(sum(dp[n])% mod)