[백준/Python] 9184_신나는 함수 실행

문제

https://www.acmicpc.net/problem/9184

9184번: 신나는 함수 실행

풀이

DP를 이용해서 한번 구했던 답을 다시 연산하지 않도록 합니다.

문제에서 주어진 아래의 조건에 의해 3차원 배열 dp의 크기는 (21) * (21) * (21)입니다. 

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

w 함수를 선언하고 재귀 연산 중간에 한번 구했던 답은 다시 연산하지 않도록 아래의 코드를 추가합니다.

if dp[a][b][c]:  # 0이 아닐 때
    return dp[a][b][c]

코드

import sys


def w(a, b, c):
    if a <= 0 or b <= 0 or c <= 0:
        return 1
    if a > 20 or b > 20 or c > 20:
        return w(20, 20, 20)
    if dp[a][b][c]:
        return dp[a][b][c]
    if a < b < c:
        dp[a][b][c] = w(a, b, c - 1) + w(a, b - 1, c - 1) - w(a, b - 1, c)
        return dp[a][b][c]
    dp[a][b][c] = w(a - 1, b, c) + w(a - 1, b - 1, c) + w(a - 1, b, c - 1) - w(a - 1, b - 1, c - 1)
    return dp[a][b][c]


dp = [[[0] * 21 for _ in range(21)] for _ in range(21)]
while True:
    A, B, C = map(int, sys.stdin.readline().rstrip().split())
    if A == -1 and B == -1 and C == -1:
        break
    print('w({}, {}, {}) = {}'.format(A, B, C, w(A, B, C)))