[DP] 금광

금광

(이것이 코딩테스트다 참조)

문제 설명

n*m 크기의 금광이 존재하고, 각 칸은 특정한 크기의 금이 들어있다.
채굴자는 첫번째 열의 어느행에서든 출발이 가능하다. 이후 m번에 걸쳐서 매번 오른쪽 위, 오른쪽, 오른쪽 아래 3가지 중 하나의 위치로 이동을 한다.
채굴자가 얻을 수 있는 금의 최대 크기를 출력해야한다.

입력

첫째줄은 테스트의 개수이다.

2
3 4
1 3 3 2 2 1 4 1 0 6 4 7
4 4
1 3 1 5 2 2 4 1 5 0 2 3 0 6 1 2

출력

19
16

코드

import sys
t = int(input())
for _ in range(t):
    n,m = map(int,input().split()) #세로 가로
    arr=list(map(int,input().split()))
        
    #2차원 테이블로 바꾸기
    dp=[]
    idx=0
    for i in range(n):
        dp.append(arr[idx:idx+m])
        idx+=m

    #dp, i열에서 최댓값을 결정하는 방식
    for j in range(1,m):
        for i in range(n):
            #왼쪽 아래 혹은 왼쪽만 가능
            if i==0:
                dp[i][j] += max(dp[i][j-1],dp[i+1][j-1])
                #왼쪽 위 혹은 왼쪽만 가능
            elif i==n-1:
                dp[i][j] +=max(dp[i-1][j-1],dp[i][j-1])
            else:
                dp[i][j]+=max(dp[i-1][j-1],dp[i][j-1],dp[i+1][j-1])
                        
    ans=0
    for i in range(n):
        ans=max(ans,dp[i][m-1]) 
    print(ans)         

i번째 열(i>=1)을 기준으로 생각하면 왼쪽 위, 왼쪽, 왼쪽 아래의 이동을 통해서 (i,j)의 위치로 올 수 있다.
문제를 쪼개어 i번째 열에 도착을 했을 때를 가정하여 그 때 마다 선택할 수 있는 최대값을 구하면 된다.

그렇게 하면 점화식은
$dp[i][[j]=arr[i][j]+max(dp[i-1][j-1],dp[i-1][[j],dp[i-1][j+1])$이 된다.
(0번째 행과 n-1번째 행 제외)