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

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

 

 

사용 알고리즘: 재귀, DP

 

 

사실 문제를 처음 읽고 어엉? 했지만 다시 보니 주어진 코드를 거의 그대로 사용해서 풀 수 있는 문제였다.

 

문제에서 주어진 코드는 재귀만을 사용해서 w(a, b, c) 를 구하기 때문에 시간이 너무 오래걸린다. 시간을 단축하기 위해 DP 를 사용해서 전에 계산했던 값을 가져다 쓰면서 시간을 줄였다. 

 

 

- dp 배열 초기화

 

문제에 주어진 코드를 읽어보면 a, b, c 중 한개라도 20 을 초과하는 값이 있다면 w(20, 20, 20) 로 통일해 버린다. 

그래서 n 값은 최대 50 까지 가능하지만, dp 배열은 51 까지 만들 필요없이 21 까지만 만들면 된다. 

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

 

 

- 함수 w 

 

문제에서 주어진 코드와 거의 똑같지만, 전에 계산해놨던 값이 있다면 가져다 사용해야 하므로 3 번째 if 문을 추가해 주었다.

그리고 마지막 if-else 문에서는 단순히 재귀만 이용하지 않고 계산한 값을 dp 배열에 넣어 다음에 사용할 수 있도록 해주었다.

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]: # dp 배열에 값이 존재하면 가져다 사용 -> 전에 이미 계산했던 적 있는 값
        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]
    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]

 

 

- 입력

 

-1 -1 -1 이 입력되면 break 해서 종료하고, 그 외에는 w 함수를 불러서 값을 계산하고 출력해주면 된다. 

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)}')

 

 

※ 참고 

 

파이썬 문자열 포맷팅: f-string 문법

 

오랜만에 문자열 포맷팅을 쓰자니 기억 나지 않아서 찾아봤다. f-string 문법은 python 3.6 이상의 버전부터 사용할 수 있다. 

사용법은 간단하다. 문자열 앞에 f 를 붙이고 변수는 중괄호 {} 안에 넣으면 된다. 

 

a = '백준'

print(f'나는 오늘 {a} 문제를 풀었다')
# 출력: 나는 오늘 백준 문제를 풀었다

 

 

전체 코드

import sys
input = sys.stdin.readline

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]: # dp 배열에 값이 존재하면 가져다 사용 -> 전에 이미 계산했던 적 있는 값
        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]
    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]

# 3차원 dp 배열 초기화
dp = [[[0 for _ in range(21)] for _ in range(21)] for _ in range(21)]

# 입력
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)}')