백준 2616번: 소형기관차 (Python)

접근

https://www.acmicpc.net/problem/2616

2616번: 소형기관차

누적합 문제인 것은 눈치 챘는데 재귀로 풀고나서 시간 초과 발생한 후에서야 dp 문제인 것을 깨달았습니다.

일단 연속된 객차를 끌 수 있기 때문에 누적합 배열을 만들어 길이가 L 인 객차의 승객 수를 한번의 계산으로 구할 수 있도록 계산해줍니다. 여기까지는 기본적인 누적합 문제와 같습니다.

dp 를 dp[N + 1][4] 로 초기화하고 dp[i][j] 에 i 번째 객차까지 j 개의 소형기관차를 사용한 경우 승객의 최대값을 저장하도록 dp 를 계산해줍니다. dp[i][j] 에 들어갈 값은

1. 바로 전 칸까지의 값 dp[i - 1][j]
2. dp[i - L][j - 1] + sum_train[i] - sum_train[i - L]

중 큰 값입니다. 2번항은 현재 위치에 소형 기관차를 놓는다고 생각하고, i - L ~ i 칸의 승객수 + 소형기관차 1개 적은 승객 최대값을 더해준다는 의미입니다.

코드

import sys

input = sys.stdin.readline

N = int(input())
train = list(map(int, input().split(" ")))
L = int(input())
ans = 0

sum_train = [0] * (N + 1)
for i in range(1, N + 1):
    sum_train[i] = sum_train[i - 1] + train[i - 1]
dp = [[0] * (4) for _ in range(N + 1)]

for j in range(1, 4):
    for i in range(1, N + 1):
        if i - L * j < 0 or N + 1 - L * (3 - j) <= i:
            continue
        dp[i][j] = max(dp[i - 1][j], dp[i - L][j - 1] + sum_train[i] - sum_train[i - L])

print(dp[-1][-1])

더 생각해 볼 것?

...

코드나 내용 관련 조언, 부족한 점 및 질문 언제든 환영합니다!