[백준] 9184 : 신나는 함수 실행 (Python/파이썬)

문제 링크

https://www.acmicpc.net/problem/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인 경우는 입력의 마지막을 제외하면 없다.

출력

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

예제

입력

1 1 1
2 2 2
10 4 6
50 50 50
-1 7 18
-1 -1 -1

출력

w(1, 1, 1) = 2
w(2, 2, 2) = 4
w(10, 4, 6) = 523
w(50, 50, 50) = 1048576
w(-1, 7, 18) = 1

---

정답 및 풀이

import sys
input = sys.stdin.readline

dp = [[[0] * 21 for _ in range(21)] for _ in range(21)]

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] != 0:
        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)
    else:
        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 == -1 and b == -1 and c == -1:
        break
    print(f"w({a}, {b}, {c}) = {w(a, b, c)}")

1. import sys
   • 파이썬 표준 입출력 모듈을 사용하기 위해 불러온다
2. input = sys.stdin.readline
   • 입력 속도를 높이기 위해 기본 input 대신 readline을 사용한다
   • 재귀 호출이 많은 문제에서 입출력 지연을 줄이기 위함이다
3. dp = [[[0] * 21 for _ in range(21)] for _ in range(21)]
   • w(a,b,c)의 결과를 저장하기 위한 3차원 dp 배열이다
   • a,b,c의 유효 범위가 0~20 이므로 크기를 21로 잡는다
   • 초기값 0은 아직 계산되지 않은 상태를 의미한다
4. def w(a, b, c):
   • 문제에서 정의한 함수 w(a,b,c)를 그대로 구현한 재귀 함수다
5. if a <= 0 or b <= 0 or c <= 0:
   • 문제의 기저 조건이다
   • 셋 중 하나라도 0 이하이면 더 이상 재귀를 진행하지 않는다
6. return 1
   • 기저 조건에 해당하므로 결과를 1로 확정한다
7. if a > 20 or b > 20 or c > 20:
   • 입력 값이 정의된 범위를 초과하는 경우를 처리한다
8. return w(20, 20, 20)
   • 20을 초과하는 값은 모두 w(20,20,20)으로 치환한다
9. if dp[a][b][c] != 0:
   • 해당 상태가 이미 계산된 적이 있는지 확인한다
   • 중복 재귀 호출을 방지하기 위한 메모이제이션 검사다
10. return dp[a][b][c]
    • 이전에 계산된 값을 바로 반환한다
11. if a < b < c:
    • 문제에서 주어진 특수한 조건 분기다
    • a,b,c가 엄격히 증가하는 경우에만 적용된다
12. dp[a][b][c] = w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, c)
    • a는 유지하고 b와 c를 줄이면서 값을 계산한다
    • 중복 계산을 제거하기 위해 마지막 항을 뺀다
13. else:
    • a < b < c 조건에 해당하지 않는 모든 경우다
14. dp[a][b][c] = (
    • 일반적인 점화식을 적용하기 시작한다
15. w(a-1, b, c)
    • a만 감소시킨 상태의 결과를 더한다
16. + w(a-1, b-1, c)
    • a와 b를 동시에 감소시킨 상태의 결과다
17. + w(a-1, b, c-1)
    • a와 c를 감소시킨 상태의 결과다
18. w(a-1, b-1, c-1)
    • 앞의 세 항에서 중복된 부분을 제거한다
19. return dp[a][b][c]
    • 계산된 값을 dp에 저장한 뒤 반환한다
20. while True:
    • 여러 테스트 케이스를 처리하기 위한 반복문이다
21. a, b, c = map(int, input().split())
    • 한 줄에서 세 정수를 입력받는다
22. if a == -1 and b == -1 and c == -1:
    • 입력 종료 조건을 확인한다
23. break
    • 종료 조건이 만족되면 반복을 중단한다
24. print(f"w({a}, {b}, {c}) = {w(a, b, c)}")
    • 문제에서 요구한 출력 형식에 맞춰 결과를 출력한다
    • w 함수는 dp 기반으로 빠르게 계산된다

---