[Python] 백준 2616 소형기관차

 

 

문제

기차는 맨 앞에 있는 기관차 1대가 손님이 탄 객차 여러 칸을 끌고 간다. 기관차가 고장나면 기차를 운행할 수 없게 되므로 최근 철도청은 기관차 고장에 대비하여 몇몇 역에 소형 기관차 3대를 배치하기로 결정하였다. 소형 기관차는 평소에 이용하는 기관차보다 훨씬 적은 수의 객차만을 끌 수 있다.

기관차가 고장났을 때 끌고 가던 객차 모두를 소형 기관차 3대가 나누어 끌 수 없기 때문에, 소형 기관차들이 어떤 객차들을 끌고 가는 것이 좋을까하는 문제를 고민하다가 다음과 같이 하기로 결정하였다.

1. 소형 기관차가 최대로 끌 수 있는 객차의 수를 미리 정해 놓고, 그보다 많은 수의 객차를 절대로 끌게 하지 않는다. 3대의 소형 기관차가 최대로 끌 수 있는 객차의 수는 서로 같다.
2. 소형 기관차 3대를 이용하여 최대한 많은 손님을 목적지까지 운송하도록 한다. 각 객차 마다 타고 있는 손님의 수는 미리 알고 있고, 다른 객차로 손님들이 이동하는 것은 허용하지 않는다.
3. 각 소형 기관차는 번호가 연속적으로 이어진 객차를 끌게 한다. 객차는 기관차 바로 뒤에 있는 객차부터 시작하여 1번 부터 차례로 번호가 붙어있다.

예를 들어 기관차가 끌고 가던 객차가 7칸이고, 소형 기관차 1대가 최대로 끌 수 있는 객차 수는 2칸이라고 하자. 그리고 1번 부터 7번까지 각 객차에 타고 있는 손님의 수가 아래 표와 같다고 하자. 괄호속에 있는 숫자는 객차 번호를 나타낸다.

(1)(2)(3)(4)(5)(6)(7)

35

40

50

10

30

45

60

소형 기관차 3대는 각각 1-2번, 3-4번, 그리고 6-7번 객차를 끌고 가면 손님 240명을 운송할 수 있고, 이보다 많은 수의 손님을 운송할 수 없다.

기관차가 끌고 가던 객차의 수와 각 객차에 타고 있던 손님의 수, 그리고 소형 기관차가 최대로 끌수 있는 객차의 수가 주어질 때, 소형 기관차 3대를 이용하여 최대로 운송할 수 있는 손님 수를 구하는 프로그램을 작성하시오.

 

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

 

 

 

 

 

 

 

난이도

골드 3

 

 

 

 

 

내 코드

import sys

n = int(sys.stdin.readline())
lst = [0] + list(map(int,sys.stdin.readline().split()))
m = int(sys.stdin.readline())
answer = 0

for i in range(n) :
    lst[i+1] += lst[i]

dp = [[0 for _ in range(n+1)] for _ in range(3)]

# dp를 3번 업데이트하며 답 찾기
for i in range(m,n+1) :
    tmp = lst[i] - lst[i-m]
    dp[0][i] = max(dp[0][i-1],tmp)

# 2
for i in range(2*m, n+1) :
    tmp = lst[i] - lst[i-m]
    dp[1][i] = max(dp[1][i-1],tmp + dp[0][i-m])

# 3
for i in range(3*m, n+1) :
    tmp = lst[i] - lst[i-m]
    dp[2][i] = max(dp[2][i-1], tmp + dp[1][i-m])

print(dp[2][n])

 

 

 

풀이과정

이 문제는 누적합과 DP로 풀 수 있다.

우선 lst배열을 누적합 배열로 만들었다. 누적합 배열을 사용하면 구간 합을 구하기 쉽고 빨라진다.

그리고 3번의 dp를 이용해서 답을 구하였는데 원리를 살펴보면,

 

첫 번째 dp 업데이트는 현재 index에서 앞에 m개의 index를 더한 값을 tmp에 저장하여두고 그것과 dp[i-1] 값을 비교하여 현재 인덱스까지의 m개씩 더했을 때의 최댓값들을 찾아두었다.

 

두 번째 dp 업데이트는 똑같이 tmp를 계산해두고, 그것에 첫 번째에 m개를 더한 값의 최댓값인 dp[0][i-m]을 더해서 dp[i-1] 값과 비교하여 2개의 기관차를 m개씩 더했을 때의 최대값을 찾았고

세 번째는 두 번째와 마찬가지로 진행하였다.

 

dp를 계산하는데 조금 시간이 걸린 문제였다.