[1309] 동물원 - Python

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

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

dp 문제다. 앞 칸에 사자가 왼쪽에 있는지, 오른쪽에 있는지, 한 마리도 없는지에 따라 다음 칸에 사자를 배치할 수 있는 경우의 수가 달라진다.

 

처음에는 [[1,1,1], [2,3,2] ... ]와 같은 dp 리스트를 만들어 각 칸에 왼쪽, 0마리, 오른쪽의 가짓수를 저장하여 풀었다. dp[1]이 [2,3,2]라는 것은 두 번째 줄에 왼쪽에 사자가 있기까지의 경우의 수는 2개, 사자가 0마리이기 까지의 경우의 수는 3개, 오른쪽에 사자가 있기까지의 경우의 수는 2개라는 뜻이다.

# 메모리 초과 코드
n = int(input())
dp = [[0, 0, 0] for _ in range(n)]
dp[0] = [1, 1, 1]
for i in range(1, n):
    dp[i][0] = dp[i-1][1] + dp[i-1][2]
    dp[i][1] = dp[i-1][0] + dp[i-1][1] + dp[i-1][2]
    dp[i][2] = dp[i-1][0] + dp[i-1][1]
print(sum(dp[n-1]) % 9901)

그러나 위와 같이 코드를 작성하면 메모리 초과가 나온다. dp[i][0]과 dp[i][2]는 항상 같다(대칭형)는 것을 이용하여 dp 리스트의 크기를 아래와 같이 줄여 보았으나 역시 메모리 초과가 떴다.

# 메모리 초과 코드
n = int(input())
dp = [[0, 0] for _ in range(n)]
dp[0] = [1, 1]
for i in range(1, n):
    dp[i][0] = dp[i-1][0] + 2 * dp[i-1][1]
    dp[i][1] = dp[i-1][0] + dp[i-1][1]
print((dp[n-1][0] + 2*dp[n-1][1]) % 9901)

 

그래서 생각한 것은 더이상 필요하지 않은 메모리는 바로 삭제하는 것이다. 각 줄의 경우의 수를 계산할 때 필요한 것은 앞 줄의 경우의 수 뿐이다. 따라서 한 줄을 계산하고 나면 그 값을 바로 dp에 덮어써서 메모리를 재사용한다.

# 정답 코드
n = int(input())
dp = [1, 1]
for _ in range(1, n):
    a = dp[0] + 2 * dp[1]
    b = dp[0] + dp[1]
    dp = [a, b]
print((dp[0] + 2*dp[1]) % 9901)