[백준] 9184 : 신나는 함수 실행 - Python

신나는 함수 실행

재귀 호출만 생각하면 신이 난다! 아닌가요?
다음과 같은 재귀함수 w(a, b, c)가 있다.

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)를 출력하는 프로그램을 작성하시오.

입력

입력은 세 정수 a, b, c로 이루어져 있으며, 한 줄에 하나씩 주어진다. 입력의 마지막은 -1 -1 -1로 나타내며, 세 정수가 모두 -1인 경우는 입력의 마지막을 제외하면 없다.

출력

입력으로 주어진 각각의 a, b, c에 대해서, w(a, b, c)를 출력한다.

import sys
sys.stdin = open("input.text", "rt")

dp = [[[0] * 21 for _ in range(21)] for _ in range(21)]
dp[0][0][0] = 1 #자명
def DFS(a,b,c):
    if a <= 0 or b <= 0 or c <= 0:
        return dp[0][0][0]
    elif a > 20 or b > 20 or c > 20:
        return DFS(20,20,20)
    
    if dp[a][b][c] > 0: #이미 값 존재
        return dp[a][b][c]
    #만약 기존 dp에 저장되어 있지 않다면 계속 해줘야함.
    
    if a<b and b<c:
        dp[a][b][c] = DFS(a,b,c-1) + DFS(a,b-1,c-1) - DFS(a,b-1,c)
    else:
        dp[a][b][c] = DFS(a-1,b,c) + DFS(a-1,b-1,c) + DFS(a-1,b,c-1) - DFS(a-1,b-1,c-1)
    
    return dp[a][b][c]
    


while True:
    a,b,c = map(int, input().split())
    if a == -1 and b == -1 and c == -1:
        break
    res = DFS(a,b,c)
    print(f"w({a}, {b}, {c}) = {res}")

코멘트

해당 문제를 풀 때는 연산 시간을 생각해서 기존의 값에 저장해둬야 한다.
그렇기에 먼저 return 하기 전에 dp 테이블에 저장을 해두고 마지막에 return 하는 것이 맞다.
그냥 문제에 있는 그대로 함수 구현 후에, dp 값이 존재하면 리턴.