[골드 3] 백준 2616 - 소형기관차 (파이썬)

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

풀이

• 열차는 각 객차마자 탑승 인원이 주어져있다.
• 한번에 K개의 객차만을 끌 수 있다.
• 최대 3번까지 소형 기관차 운행이 가능하다.
• 최대한 많은 승객을 운송하는 것이 목표이다.

먼저 입력 및 누적합을 계산한다.

N = int(input())  # 객차 수
A = list(map(int, input().split()))  # 각 객차에 탄 승객 수
K = int(input())  # 소형 기관차가 끌 수 있는 객차 수

psum = [0] * N
psum[0] = A[0]
for i in range(1, N):
    psum[i] = psum[i - 1] + A[i]

• psum[i]는 0부터 i번째까지의 누적합이다.
• 구간 합 A[i - K + 1] … + A[i]는 psum[i] - psum[i - K]로 계산한다.

DP 테이블은 아래와 같이 정의한다.

D = [[0] * 4 for _ in range(N)]

• D[i][j]는 0 ~ i 번째 까지 객차 중 j번 기관차를 사용했을때의 최대 승객 수를 의미한다.
• j는 1 ~ 3까지 사용한다.

점화식은 아래와 같이 구성한다.

for i in range(1, 4):
    D[0][i] = A[0]  # 첫 객차만 사용할 경우 초기값 (실제론 큰 의미 없음)

for i in range(1, N):
    for j in range(1, 4):
        D[i][j] = D[i - 1][j]  # 현재 객차 포함하지 않는 경우

        if i <= K - 1:
            D[i][j] = max(D[i][j], psum[i])  # 처음부터 i까지의 합을 고를 수 있는 경우
        else:
            sum = psum[i] - psum[i - K]
            D[i][j] = max(D[i][j], sum + D[i - K][j - 1])

• D[i][j] = max(D[i - 1][j], sum(i - K + 1 ~ i) + D[i - K][j - 1])
• 현재 구간을 고르지 않는 경우는 D[i - 1][j]
• 현재 구간을 고르는 경우는 이전 구간까지의 최대합 + 현재 구간 합을 지정한다.

psum은 N + 1 길이로 만들고 psum[i] = sum(A[0] ~ A[i - 1]) 식으로 하면 인덱싱이 깔끔해 질 것 같다.

코드

# 백준 2616 - 소형기관차
# 분류 - 다이나믹 프로그래밍

N = int(input())
A = list(map(int, input().split()))
K = int(input())

# sum(i - K + 1, i) = psum[i] - psum[i - K]
psum = [0] * N
psum[0] = A[0]
for i in range(1, N) :
    psum[i] = psum[i - 1] + A[i]

D = [[0] * 4 for _ in range(N)]

for i in range(1, 4) :
    D[0][i] = A[0]

for i in range(1, N) :
    for j in range(1, 4) :
        D[i][j] = D[i - 1][j]

        if i <= K - 1 :
            D[i][j] = max(D[i][j], psum[i])
        else :
            sum = psum[i] - psum[i - K]
            D[i][j] = max(D[i][j], sum + D[i - K][j - 1])

print(D[N - 1][3])