해커를 꿈꾸는 개발자

파이썬 / BOJ 백준 / 9184번 신나는 함수 실행 - dp

문제

if a <= 0 or b <= 0 or c <= 0, then w(a, b, c) returns:
  1

if a > 20 or b > 20 or c > 20, then w(a, b, c) returns:
  w(20, 20, 20)

if a < b and b < c, then w(a, b, c) returns:
  w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, c)

otherwise it returns: w(a-1, b, c) + w(a-1, b-1, c) + w(a-1, b, c-1) - w(a-1, b-1, c-1)

위의 함수를 구현하는 것은 매우 쉽다. 하지만, 그대로 구현하면 값을 구하는데 매우 오랜 시간이 걸린다. (예를 들면, a=15, b=15, c=15)
a, b, c가 주어졌을 때, w(a, b, c)를 출력하는 프로그램을 작성하시오.

풀이

1. 문제에서 주어진 함수를 그대로 사용하면, 시간초과가 발생합니다.
2. 재귀함수를 사용할 때, 동일한 계산을 수행하지 않도록 메모이제이션을 사용하는 방법을 묻는 문제입니다.
3. 계산한 값을 dp[a][b][c]에 저장하고, 이후 a, b, c를 구할 때, 계산하지 않고, 저장한 dp[a][b][c]를 return하게 되면 수행시간이 기하급수적으로 단축됩니다.

전체코드

import sys

def w(a1,b1,c1):
    if dp[a1][b1][c1] != -1:
        return dp[a1][b1][c1]

    if a1 <= 0 or b1 <= 0 or c1 <= 0:
        dp[a1][b1][c1] = 1
        return dp[a1][b1][c1]
    if a1 > 20 or b1 > 20 or c1 > 20:
        dp[a1][b1][c1] = w(20, 20, 20)
        return dp[a1][b1][c1]
    if a1 < b1 and b1 < c1:
        dp[a1][b1][c1] = w(a1, b1, c1-1) + w(a1, b1-1, c1-1) - w(a1, b1-1, c1)
        return dp[a1][b1][c1]
    else :
        dp[a1][b1][c1] = w(a1-1, b1, c1) + w(a1-1, b1-1, c1) + w(a1-1, b1, c1-1) - w(a1-1, b1-1, c1-1)
        return dp[a1][b1][c1]

    pass

dp = [[[-1 for _ in range(n)] for _ in range(n)] for _ in range(n)]
while True:
    a,b,c = map(int, sys.stdin.readline().split())

    if a == -1 and b == -1 and c == -1:
        break

    result = w(a,b,c)
    p = f'w({a}, {b}, {c}) = {result}'
    print(p)