[백준] 2616 : 소형기관차 (파이썬)

DP, 누적합 문제

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

 

이 문제는 냅색 알고리즘의 변형으로도 볼 수 있을 것 같다.

누적합을 저장해준 (1차원) 배열을 만들고, 해당 배열을 활용하여 소형 기관차 3대로 최대 운송할 수 있는 손님수를 구하는 문제이다.

 

초기 접근 방식

• 객차의 수, 소형 기관차가 최대로 끌 수 있는 객차의 수, 소형 기관차의 갯수(3)에 대한 정보가 주어졌다.
  처음에 이 문제를 보고 DP로 풀어야겠다고 생각했다. 그러나 초기에는 r = 소형 기관차가 최대로 끌 수 있는 객차의 수, c = 객차의 수에 대한 DP 배열의 형태를 구상하였기에 문제가 생각보다 쉽게 풀리지 않았다.
  
  3대의 소형 기관차가 최대로 끌 수 있는 객차의 수는 서로 같다. 여기서 "최대"라는 말에 포커스가 맞춰져서 최대일 경우와 최대가 아닐 경우로 구분하여 생각했던 것 같다.
• 또, 처음에는 누적합 배열 대신 구간 <=  소형 기관차가 최대로 끌 수 있는 객차의 수를 갖는 구간합 배열을 생각했다. 하지만, 이러한 구간합 배열을 활용할 경우 백트래킹과 같은 방식으로 경우의 수를 조합해야 한다.
  최대 객차의 수는 50000으로 구간합 + 백트래킹 방식을 사용할 경우 50000C3에 해당하는 경우의 수를 훑어야 한다.

 

💡 풀이코드 (성공)

import sys 
n = int(sys.stdin.readline())
arr = [0]
l = list(map(int, sys.stdin.readline().split()))
arr += l
limit = int(sys.stdin.readline())

# 누적합 구하기 
prefix_sum = [0] * (n + 1)
for i in range(1, n + 1):
    prefix_sum[i] = prefix_sum[i - 1] + arr[i]

def sol(arr, n, prefix_sum, limit):
    dp = [[0] * (n + 1) for _ in range(4)]

    for i in range(1, 4):
        for j in range(n + 1):
            dp[i][j] = max(dp[i][j - 1], dp[i-1][j-limit] + prefix_sum[j] - prefix_sum[j-limit])
    print(dp[3][n])

sol(arr, n, prefix_sum, limit)

 

1. 누적합 배열

반복되는 덧셈 연산을 피하기 위해 누적합 배열을 만들었다.

prefix_sum = [0] * (n + 1)
for i in range(1, n + 1):
    prefix_sum[i] = prefix_sum[i - 1] + arr[i]

 

2. DP 배열 

• r = 4(소형기관차 갯수 + 1), c = n(객차의 수 + 1)의 2차원 DP 배열을 생성하였다.
• dp[i][j] = i개의 소형 기관차가 j개의 객차까지 고려했을 때 운송할 수 있는 최대 승객의 수

 

3. 점화식

dp[i][j] = max(dp[i][j - 1], dp[i-1][j-limit] + prefix_sum[j] - prefix_sum[j-limit])

 

dp[i][j]
= max(현재 i번째 소형기관차가 j번째 객차를 추가하지 않는 경우, 현재 i번째 소형기관차가 j번째 객차를 추가하는 경우)
= max(현재 i번째 소형기관차가 j-1번째까지 고려했을 때 운송할 수 있는 최대 승객의 수, 이전 i-1번째 소형기관차가 j-limit 객차까지 고려했을 때 운송할 수 있는 최대 승객의 수 + (j-limit ~ j)번째까지의 승객의 수)

= max(dp[i][j - 1], dp[i-1][j-limit] + prefix_sum[j] - prefix_sum[j-limit])

 

4. 출력값

 

dp[3][n] = 3개의 소형 기관차가 n번째(마지막) 객차까지 고려했을 때 운송할 수 있는 최대 승객의 수

 

느낀점

DP 문제는 어떠한 정보를 활용하여 DP 배열을 만들 것인지가 정말 중요한 것 같다.

 

 

참고문헌

https://devjoy.tistory.com/265

[백준] 2616 : 소형기관차 / python

https://taxol1203.github.io/codingtest/bj-%EC%86%8C%ED%98%95%EA%B8%B0%EA%B4%80%EC%B0%A8/

백준 2616 - 소형기관차