[BOJ][⚪2][백준#09184] 신나는 함수 실행

문제 출처

BAEKJOON Online Judge #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)를 출력한다.

제한

-50 ≤ a, b, c ≤ 50

예제

입력

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

My Sol

import sys
input = sys.stdin.readline

def make(a, b, c):
    if memo[a][b][c]!=-1: pass
    elif a>70 or b>70 or c>70:
        memo[a][b][c] = make(70, 70, 70)
    elif a < b < c:
        memo[a][b][c] = make(a, b, c-1) + make(a, b-1, c-1) - make(a, b-1, c)
    else:
        memo[a][b][c] = make(a-1, b, c) + make(a-1, b-1, c) + make(a-1, b, c-1) - make(a-1, b-1, c-1)
    return memo[a][b][c]


memo = [[[1]*101 for _ in range(101)] for _ in range(101)]
for x in range(51, 101):
    for y in range(51, 101):
        for z in range(51, 101):
            memo[x][y][z] = -1


while True:
    a, b, c, = map(lambda x: int(x)+50, input().split())
    if a==b==c==49: break
    print(f'w({a-50}, {b-50}, {c-50}) = {make(a, b, c)}')

메모이제이션, DP를 사용할 수 있는지에 대한 문제였다. 메모이제이션 배열을 만들어 DP를 할 줄만 안다면 어떤 방식이든 쉽게 풀어낼 수 있는 문제였다.

결과

맞았습니다!!

모범답안

출처

import sys

def w(a,b,c):
    if a <= 0 or b <= 0 or c <= 0:
        return 1
    elif a > 20 or b > 20 or c > 20:
        return w(20,20,20)
    elif dp[a][b][c]:
        return dp[a][b][c]
    elif 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]

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

inputs = sys.stdin.read().splitlines()
inputs.pop()
for n in inputs:
    a,b,c = map(int,n.split())
    sys.stdout.write(f'w({a}, {b}, {c}) = {w(a,b,c)}\n')

음수 부분은 memo 없이 바로 1을 return하는 방식을 이용하여 배열의 크기를 기존 100x100x100에서 20x20x20 수준으로 1/120 만큼으로 줄인 것이 연산 시간과 메모리를 줄이는 데에 큰 기여를 했을 것이다. 뭐 물론 38000KB와 30000KB의 메모리 차이, 150ms와 100ms의 연산시간 차이이지만 말이다.

아무튼 문제에서 제시한 그대로 메모이제이션에 활용하여 손쉽게 나타내었다. 특히나 sys.stdout.write를 이용해 출력한 것이 테스트케이스가 많은 경우 출력이 빨랐을 것이다.