백준 1309번. 동물원 (Python / 파이썬)

문제

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

이 동물원에는 사자들이 살고 있는데 사자들을 우리에 가둘 때, 가로로도 세로로도 붙어 있게 배치할 수는 없다. 이 동물원 조련사는 사자들의 배치 문제 때문에 골머리를 앓고 있다.
동물원 조련사의 머리가 아프지 않도록 우리가 2*N 배열에 사자를 배치하는 경우의 수가 몇 가지인지를 알아내는 프로그램을 작성해 주도록 하자. 사자를 한 마리도 배치하지 않는 경우도 하나의 경우의 수로 친다고 가정한다.

입력

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

출력

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

---

문제 풀이 

사자들을 우리에 가둘 때 가로로도 세로로도 붙어 있게 배치할 수 없으므로 

사자들을 우리에 가두는 경우는 총 세가지 종류로 분류할 수 있다. 

1. N-1 칸에 사자가 한 마리도 배치되지 않은 경우 

2. N-1 칸에 사자가 왼쪽에 배치된 경우 

3. N-1 칸에 사자가 오른쪽에 배치된 경우 

이처럼 N-1 칸에서 사자가 어떻게 배치되어있느냐에 따라 N칸의 사자의 배치가 결정된다. 

 

N=1 일 때의 경우의 수를 생각해보면 아래 그림과 같이 

사자가 한마리도 배치되지 않은 경우, 사자가 왼쪽에 배치된 경우, 사자가 오른쪽에 배치된 경우 각각 하나씩 총 세가지 이다.  

 

N=2 부터 나머지 경우의 수는 아래와 같이 DP를 이용해서 구할 수 있다. 

0번째 칸은 세 종류의 경우의 수 중 1번을 의미하고

1번째 칸은 2번, 2번째 칸은 3번을 의미한다. 

따라서, 0번째 칸은 N-1의 모든 칸의 수를 더해주고 

1번째 칸은 N-1에서 왼쪽 칸을 제외하고 더해주고 

2번째 칸은 N-1에서 오른쪽 칸을 제외하고 더해주면 위와 같이 결과를 얻을 수 있다. 

 

총 경우의 수는 N번째 칸에 저장된 모든 경우의 수를 더하여 구할 수 있다. 

주의해야 할 점은, 문제에서 9901로 나눈 나머지를 출력하라 했는데 9901로 나누지 않으면 메모리 초과 문제가 발생한다. 따라서 꼭 9901로 나누어주어 결과를 출력해야 한다. 

---

My Code

n = int(input())

dp = [[0] * 3 for _ in range(n+1)]
dp[1][0] = 1
dp[1][1] = 1
dp[1][2] = 1 

for i in range(2, n + 1):
    for j in range(3):
        if j == 0:
            dp[i][j] = (dp[i-1][j] + dp[i-1][j+1] + dp[i-1][j+2]) % 9901
        elif j == 1:
            dp[i][j] = (dp[i-1][j-1] + dp[i-1][j+1]) % 9901
        else:
            dp[i][j] = (dp[i-1][j-1] + dp[i-1][j-2]) % 9901

print(sum(dp[n]) % 9901)