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


문제


다음과 같은 재귀함수 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




결과






댓글 쓰기

2 댓글

  1. 다음주까지 dp 관련된 백준 문제 풀어야하는데 좋은 참조가 되겠군요 감사합니다.

    답글삭제
    답글
    1. 도움이 되었다니 기쁘네요 :)

      삭제