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

백준 문제 주소 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인 경우는 입력의 마지막을 제외하면 없다.

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

출력

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

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

제한

-50 ≤ a, b, c ≤ 50

---

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 a < b < c:
        return w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, c)
    
    return 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, b, c = 0, 0, 0
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)}')

문제의 예시를 그대로 구현하면 위와 같은 코드가 나온다. 하지만 이 코드는 값을 구하는데 오랜 시간이 걸리고, 백준 기준으로는 시간 초과가 나온다.

문제의 기준을 통과하기 위해서는 문제를 작게 쪼개서 겹치는 부분의 계산을 줄이는 것이 필요하다. 이를 위해서 배열을 만들어서 값을 저장하고 필요한 부분에서 사용하였다.

Python 풀이

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 array[a][b][c]:
        return array[a][b][c]
    if a < b < c:
        array[a][b][c] =  w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, c)
        return array[a][b][c]
    
    array[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 array[a][b][c]

array = [[[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)}')

a, b, c를 0부터 20까지 값을 넣을 배열이 필요해서 3차원 배열을 선언하고, 함수에서 값을 계산하면 배열에 저장하여 같은 계산을 반복하는 것 없앴다.

Kotlin 풀이

var array = Array(21){Array(21){Array(21){0}}}

fun main(args: Array<String>)
{
    while(true) {
        val n = readLine()
        if(n == "-1 -1 -1") break
        val abc = n?.split(" ")!!.map { it.toInt() }.toTypedArray()
        println("w(${abc[0]}, ${abc[1]}, ${abc[2]}) = ${w(abc[0], abc[1], abc[2])}")
    }
}

fun w(a: Int, b: Int, c: Int): Int {
    if(a <= 0 || b <= 0 || c <= 0) {
        return 1
    }
    if(a > 20 || b > 20 || c > 20) {
        return w(20, 20, 20)
    }
    if(array[a][b][c] != 0) {
        return array[a][b][c]
    }
    if(a < b && b < c) {
        array[a][b][c] =  w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, c)
        return array[a][b][c]
    }

    array[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 array[a][b][c]
}

kotlin을 사용하여 입출력을 넣는 것이 익숙하지 않아 형식을 틀렸더니 런타임에러가 났다. 다음에 한번 날잡고 입출력 방법들에 대해서 알아봐야할 것 같다.