[BOJ] 9184번 신나는 함수 실행 - 파이썬

📒 문제

재귀 호출만 생각하면 신이 난다! 아닌가요?

다음과 같은 재귀함수 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인 경우는 입력의 마지막을 제외하면 없다.

-50 ≤ a, b, c ≤ 50

출력

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

✏️ 풀이

처음에는 문제 이해가 잘 안가서 무슨 말이지..싶었는데, 쭉 보다보니 그냥 w라는 함수를 만드는 게 문제였다. 하지만 주어진 if문을 그대로 입력하면 시간이 오래걸리기 때문에 수정해야 했다.
재귀 사용을 피하려면 dp를 써야한다고 생각하고 있었기 때문에 dp를 사용했다.
또한 if문의 조건을 보면 a, b, c는 0부터 20까지의 숫자만 의미가 있는 것이고 나머지 범위의 숫자들은 의미가 없다고 판단할 수 있기 때문에 다음과 같이 코드를 구성해보았다.

MAX = 21  # 0부터 20까지
dp = [[[0]*MAX for _ in range(MAX)] for __ in range(MAX)]

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]

while True:
    a, b, c = map(int, input().split())
    if [a, b, c] == [-1, -1, -1]:
        break
    print(f'w({a}, {b}, {c}) = {w(a, b, c)}')

입력에 따라 출력해주어야 하기 때문에 while문을 사용하여 입출력을 다루어 보았다.
메모리는 29452kb, 시간은 988ms로 적당한 결과가 나온 것 같다.