2616. 소형기관차 (Python)

2616번: 소형기관차

 

풀이

 구간합 개념과 dp를 사용하여 해결할 수 있다. dp 배열은 2차원 배열로 이루어져 있는데 의미는 다음과 같다.

    # dp[1][n]: n번째 객차까지 1개의 소형 기관차를 사용했을 때 운송할 수 있는 최대 손님 수
    # dp[2][n]: n번째 객차까지 2개의 소형 기관차를 사용했을 때 운송할 수 있는 최대 손님 수
    # dp[3][n]: n번째 객차까지 3개의 소형 기관차를 사용했을 때 운송할 수 있는 최대 손님 수
    dp = [[0 for _ in range(N+1)] for _ in range(4)]

 

 소형 기관차가 최대로 끌 수 있는 객차의 수를 M이라 하고 먼저 처음 M개의 객차까지의 손님에 대하여 합을 구한다. 이 구간합은 처음 M개의 객차까지 1개의 소형 기관차를 사용했을 때의 최댓값을 의미하므로 dp[1][M]에 구간합을 먼저 대입한다. 그 후 bottom-up 방식으로 dp 배열을 채워나간다.

 

 M+1 이후의 객차(i 번째)를 살펴보면서 해당 범위에 맞춰 구간합을 갱신한다. 그 후 dp[1][i]는 이전 i-1번째 까지의 소형 기관차 1개를 사용했을 때의 최댓값과 현재 구간합을 비교하여 갱신해나간다. 이를 통해 dp[1][i]의 경우는 i번째까지 1개의 소형 기관차를 사용했을 때 태울 수 있는 손님의 최댓값이 담기게 된다.

dp[1][i] = max(dp[1][i-1], sums) # i번째 객차까지 1개의 소형 기관차를 사용했을 때 최댓값 갱신

 

 2개(또는 3개)의 소형 기관차를 사용하는 경우는 i-1번째 객차까지 2개(또는 3개)의 소형 기관차를 사용했을 때의 태울 수 있는 손님의 최댓값과 i-M번째 객차까지 1개(또는 2개)의 소형 기관차를 사용했을 때 태울 수 있는 손님의 최댓값에 현재 구간합을 합한 값을 비교하여 큰 값을 dp[2][i](또는 dp[3][i])에 갱신해나간다. 이를 통해 dp[2][i](또는 dp[3][i])에는 i번째까지 2개(또는 3개)의 소형 기관차를 사용했을 때 태울 수 있는 손님의 최댓값이 담기게 된다.

 

 마지막에 dp[3][N]을 출력하여 답을 구할 수 있다.

 

코드

import sys
input = sys.stdin.readline

if __name__ == '__main__':
    N = int(input())  # 기관차가 끌고 가던 객차의 수
    people = [0] + list(map(int, input().split()))
    M = int(input())  # 소형 기관차가 최대로 끌 수 있는 객차의 수

    # dp[1][n]: n번째 객차까지 1개의 소형 기관차를 사용했을 때 운송할 수 있는 최대 손님 수
    # dp[2][n]: n번째 객차까지 2개의 소형 기관차를 사용했을 때 운송할 수 있는 최대 손님 수
    # dp[3][n]: n번째 객차까지 3개의 소형 기관차를 사용했을 때 운송할 수 있는 최대 손님 수
    dp = [[0 for _ in range(N+1)] for _ in range(4)] 
    sums = 0 # 구간합

    # 최초 M개의 객차까지의 구간합을 구해놓는다.
    for i in range(1, M+1): 
        sums += people[i]
    dp[1][M] = sums # M번째 객차까지는 1개의 소형 기관차만 사용할 수 있다. 

    # 이제 M번째 이후의 객차부터 차례대로 살펴본다.
    for i in range(M+1, N+1):
        sums = sums - people[i-M] + people[i] # 구간합 갱신 
        dp[1][i] = max(dp[1][i-1], sums) # i번째 객차까지 1개의 소형 기관차를 사용했을 때 최댓값 갱신

        if i >= M*2: # 2개의 소형 기관차를 채울 수 있다면
            # i-1번째 객차까지 2개의 소형 기관차를 사용했을 때의 손님 수와
            # 현재 구간합에 i-M번째 객차까지 1개의 소형 기관차를 사용했을 때의 최대 손님수를 합한 수를 
            # 비교하여 큰 값으로 갱신
            dp[2][i] = max(dp[2][i-1], sums + dp[1][i-M]) 

        if i >= M*3: # 3개의 소형 기관차를 채울 수 있다면
            # i-1번째 객차까지 3개의 소형 기관차를 사용했을 때의 손님 수와
            # 현재 구간합에 i-M번째 객차까지 2개의 소형 기관차를 사용했을 때의 최대 손님수를 합한 수를 
            # 비교하여 큰 값으로 갱신
            dp[3][i] = max(dp[3][i-1], sums + dp[2][i-M])

    print(dp[3][N])