동주의 세상

Clean and minimal personal blog

백준 2616 소형기관차 - 파이썬

Posted at — Jun 16, 2021

접근

  1. 열 길이가 기차의 수 + 1이고, 행 길이가 3인 dp 행렬로 접근한다.
  2. 1~3번째 행은 각각 현재 인덱스에서 1~3번째 열차가 골라졌다면 실을 수 있는 최대 손님의 수다.
  3. 입력으로 받은 train 리스트를 탐색하며 window 만큼 떨어져 있는 dp 를 업데이트 한다.
  4. dp 업데이트는 dp[현재 열 - window][현재행 - 1] 과 dp[현재 열 - 1][현재 행] 중 큰 것을 고른다.

사고

반복되는 연산을 최소화하는 DP 문제를 연습하는 중이니 다른 최적화도 가능하다면 시도해보자. 문제에서 소형 기관차의 길이가 길다면 겹치는 구간을 반복해서 구해야 할 것이다.

따라서 15-19 줄로 각각의 인덱스에서 시작하는 손님합을 최적화하여 미리 구해놓았다.

 

import sys
input = sys.stdin.readline

N = int(input())
trains = list(map(int, input().split()))
window = int(input())

dp = [[-1 for i in range(N + 1)] for j in range(3)]

sum_trains = [0] * (N - window + 1)
sm = 0
for i in range(window):
    sm += trains[i]

sum_trains[0] = sm
for i in range(1, N - window + 1):
    sm -= trains[i - 1]
    sm += trains[i + window - 1]
    sum_trains[i] = sm

for i in range(0, N + 1 - window):
    # 현재 인덱스에서 첫번째 열차가 시작되는 경우
    dp[0][i + window] = max(sum_trains[i], dp[0][i + window - 1])

    # 현재 인덱스에서 두번째 열차가 시작되는 경우
    if dp[0][i] > -1:
        dp[1][i + window] = max(dp[0][i] + sum_trains[i], dp[1][i + window - 1])

    # 현재 인덱스에서 세번째 열차가 시작되는 경우
    if dp[1][i] > -1:
        dp[2][i + window] = max(dp[1][i] + sum_trains[i], dp[2][i + window - 1])

print(dp[2][N])

submission