본문 바로가기

Develop/코딩인터뷰

[코딩인터뷰] 백준 2616번 소형기관차

반응형

백준 1616번 소형기관차
백준 1616번 소형기관차

DP & 누적 합

처음엔 N 개 중 간격 M 을 두고 3 개를 뽑는 모든 경우의 수를 검토하는 방법으로 풀었더니 '시간 초과' 발생했고 이 글을 참고해서 다시 풀었습니다.

핵심 코드는 dp[x][y] = max(dp[x][y-1], dp[x-1][y-M] + (s[y] - s[y-M])) 로 매 계산마다 이전 최대값과 비교하기 때문에 마지막에 답이 나오게 됩니다. 코드를 바로 이해하기 어려워서 디버거 변수를 노트에 적어가며 해보다 보니 이해하게 되었고 아래 코드로 정답을 맞출 수 있었습니다.

import sys
readline = sys.stdin.readline
read = sys.stdin.read

N = int(readline())
T = [0]
for t in map(int, readline().split()):
    T.append(t)

M = int(readline())

s = [0 for _ in range(N + 1)]
for i in range(1, N+1):
    s[i] = s[i-1] + T[i]

dp = [[0] * (N+1) for _ in range(4)]
for x in range(1, 4):
    for y in range(x * M, N + 1):
        dp[x][y] = max(dp[x][y-1], dp[x-1][y-M] + (s[y] - s[y-M]))

print(dp[3][N])
반응형