[백준] 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

풀이

첫번째 풀이

dp에 관한 감을 잡지 않고 풀었던 풀이이다.

특히 아래 코드는 (a,b,c) 에 대해 dp에서 저장된 값을 찾지 않고, 저장된 값이 없을때 실행하는 w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, 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) 부분에 대해 dp에서 저장된 값을 찾고 있다.

문제의 부분

elif a<b and b<c: 
        if (a,b,c) in w_dict: 
            return w_dict[(a,b,c)]
        else: 
            w_1 = w_dict.get((a, b, c-1), None)
            w_2 = w_dict.get((a, b-1, c-1), None)
            w_3 = w_dict.get((a, b-1, c), None)
            
            if w_1 is not None: w_1_value = w_1 
            else: w_1_value = w(a,b,c-1)
            if w_2 is not None: w_2_value = w_2 
            else: w_2_value = w(a,b-1,c-1)
            if w_3 is not None: w_3_value = w_3 
            else: w_3_value = w(a,b-1,c)
            
            w_dict[(a,b,c)] = w_1_value + w_2_value - w_3_value
            return w_dict[(a,b,c)]
    
    else:
        w_1 = w_dict.get((a-1, b, c), None)
        w_2 = w_dict.get((a-1, b-1, c), None)
        w_3 = w_dict.get((a-1, b, c-1), None)
        w_4 = w_dict.get((a-1, b-1, c-1), None)

        if w_1 is not None: w_1_value = w_1
        else: w_1_value = w(a-1, b, c)

        if w_2 is not None: w_2_value = w_2
        else: w_2_value = w(a-1, b-1, c)

        if w_3 is not None: w_3_value = w_3
        else: w_3_value = w(a-1, b, c-1)

        if w_4 is not None: w_4_value = w_4
        else: w_4_value = w(a-1, b-1, c-1)

        w_dict[(a,b,c)] = w_1_value + w_2_value + w_3_value - w_4_value
        return w_dict[(a,b,c)]

전체 코드

ans_key_arr = []
w_dict = {}
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:  
        if (20,20,20) in w_dict: return w_dict[(20,20,20)]
        else:
            w_dict[(20,20,20)] = w(20,20,20)
            return w_dict[(20,20,20)]

    elif a<b and b<c: 
        if (a,b,c) in w_dict: 
            return w_dict[(a,b,c)]
        else: 
            w_1 = w_dict.get((a, b, c-1), None)
            w_2 = w_dict.get((a, b-1, c-1), None)
            w_3 = w_dict.get((a, b-1, c), None)
            
            if w_1 is not None: w_1_value = w_1 
            else: w_1_value = w(a,b,c-1)
            if w_2 is not None: w_2_value = w_2 
            else: w_2_value = w(a,b-1,c-1)
            if w_3 is not None: w_3_value = w_3 
            else: w_3_value = w(a,b-1,c)
            
            w_dict[(a,b,c)] = w_1_value + w_2_value - w_3_value
            return w_dict[(a,b,c)]
    
    else:
        w_1 = w_dict.get((a-1, b, c), None)
        w_2 = w_dict.get((a-1, b-1, c), None)
        w_3 = w_dict.get((a-1, b, c-1), None)
        w_4 = w_dict.get((a-1, b-1, c-1), None)

        if w_1 is not None: w_1_value = w_1
        else: w_1_value = w(a-1, b, c)

        if w_2 is not None: w_2_value = w_2
        else: w_2_value = w(a-1, b-1, c)

        if w_3 is not None: w_3_value = w_3
        else: w_3_value = w(a-1, b, c-1)

        if w_4 is not None: w_4_value = w_4
        else: w_4_value = w(a-1, b-1, c-1)

        w_dict[(a,b,c)] = w_1_value + w_2_value + w_3_value - w_4_value
        return w_dict[(a,b,c)]
    
while True:
    a, b, c = map(int, input().split())
    if(a==-1 and b==-1 and c==-1): break
    ans_key_arr.append((a,b,c)) 
    w_dict[(a,b,c)] = w(a,b,c)
    
for key in ans_key_arr:
    print(f'w{key} = {w_dict[key]}')

두번째 풀이

첫번째 풀이에서 간소화할 점을 찾고, 불필요한 부분을 제거하였다.

특히, a<=0 or b<=0 or c<=0 조건에서는 (a,b,c)를 키로 갖는 값을 저장할 필요없이 1만 반환하면 되고, a>20 or b>20 or c>20 조건에서도 동일하게 (a,b,c) 를 키로 가지는 값을 저장할 필요없이 w(20,20,20)의 값만 구해서 반환하면 된다.

전체 코드

ans_key_arr = []
w_dict = {}
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: 
        if (20,20,20) in w_dict: return w_dict[(20,20,20)]
        else: 
            w_dict[(20,20,20)] = w(20,20,20)
            return w_dict[(20,20,20)]
        
    elif a<b and b<c: 
        if (a,b,c) in w_dict: return w_dict[(a,b,c)]
        else:
            w_dict[(a,b,c)] = w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, c)
            return w_dict[(a,b,c)]
    
    else:
        if (a,b,c) in w_dict: return w_dict[(a,b,c)]
        else: 
            w_dict[(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 w_dict[(a,b,c)]
    
while True:
    a, b, c = map(int, input().split())
    if(a==-1 and b==-1 and c==-1): break
    ans_key_arr.append((a,b,c)) 
    w_dict[(a,b,c)] = w(a,b,c)
    
for key in ans_key_arr:
    print(f'w{key} = {w_dict[key]}')