[백준 9184번] 동적 계획법(DP) 문제를 파이써닉하게 풀어보자

문제

백준 9184번 "신나는 함수 실행"

다음과 같은 재귀함수 w(a, b, c)가 있다.

위의 함수를 구현하는 것은 매우 쉽다. 하지만, 그대로 구현하면 값을 구하는데 매우 오랜 시간이 걸린다. (예를 들면, a=15, b=15, c=15)

a, b, c가 주어졌을 때, w(a, b, c)를 출력하는 프로그램을 작성하시오.

• 입력: 세 정수 a, b, c로 이루어져 있으며, 한 줄에 하나씩 주어진다. 입력의 마지막은 -1 -1 -1로 나타내며, 세 정수가 모두 -1인 경우는 입력의 마지막을 제외하면 없다.

• 출력: w(a, b, c)를 출력한다.

• 조건: -50 ≤ a, b, c ≤ 50

풀이

1. 입출력

우선 arr, res 두 빈 배열을 만들고, arr 배열에 (a, b, c) 를 입력을 받으면 w(a, b, c) 연산을 수행해 res 에 담기로 한다. 다만 "-1, -1, -1" 이 입력될 경우 입력을 그만 받아야 하므로 조건문을 넣어준다.

while True:
    a, b, c = map(int, input().split())
    if (a, b, c) == (-1, -1, -1):
        break
    arr.append((a, b, c))
    res.append(w(a, b, c))
for i, r in enumerate(res):
    print("w{} = {}".format(arr[i], r))

출력되는 부분은 깔끔함을 여지 없이 보여주는 enumerate 함수를 사용해, 배열에서 for 문을 통해 인덱스와 값을 동시에 추출할 수 있다. 또한 " ".format 을 사용해 요구하는 형식을 단 한줄로 맞출 수 있다.

2. 동적 계획법

우선 0으로 초기화 된 3차원 배열 형태의 DP 테이블을 만든다. 이곳에 (a, b, c) 에 해당하는 w(a, b, c) 결과값이 저장될 것이다.

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

DP 테이블에서 해당 주소 (a, b, c) 에 이미 연산된 값이 존재한다면 그대로 return, 없다면 새로 연산하는 코드를 구현한다. 이 로직을 통해 재귀함수의 불필요한 중복 연산을 방지해 연산량을 줄일 수 있다.

입력되는 a, b, c 의 조건이 -50부터 시작하므로, 인덱스 0으로 시작하기 위해 각각 +50을 넣어준다.

def w(a, b, c):
		# dp 테이블에 이미 연산되어 저장된 값이 있다면, 그 값을 리턴
    if dp[a+50][b+50][c+50] != 0:
        return dp[a+50][b+50][c+50]
        
		# 조건에 맞춰 연산을 수행한 후, 그 결과값을 dp 테이블에 저장 후 리턴
    if a <=0 or b <= 0 or c <= 0:
        dp[a+50][b+50][c+50] = 1
        return 1
    if a > 20 or b > 20 or c > 20:
        temp = w(20, 20, 20)
        dp[a+50][b+50][c+50] = temp
        return temp
    if a < b and b < c:
        temp = w(a, b, c-1) + w(a, b-1, c-1) - w(a, b-1, c)
        dp[a+50][b+50][c+50] = temp
        return temp
    else:
        temp = 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[a+50][b+50][c+50] = temp
        return temp

깃허브에서 전체 코드 보기

결과