백준 2616번 소형기관차 - C++ 풀이

---

 

1. sum[i] = arr[i...i+K-1]의 구간합
2. dp(idx, n): [idx번 객차~마지막 객차]가 있을 때 n개의 소형 기관차로 끌 수 있는 최대 손님의 수
3. dp(idx, n) = max(dp(idx+1, n), sum[idx] + dp(idx+K, n-1))

 

1. sum[i] = arr[i...i+K-1]의 구간합

 먼저 모든 길이 K인 구간의 구간합을 구해둔다. 누적합을 사용하여 O(N)에 구할 수 있다. sum[i] = sum[i-1] - arr[i-1] + arr[i+K-1]

 

2. dp(idx, n): [idx번 객차~마지막 객차]가 있을 때 n개의 소형 기관차로 끌 수 있는 최대 손님의 수

 dp(idx, n)을 위와 같이 정의한다. 총 구해야 하는 상태 값은 50,000 * 3개.

 

3. dp(idx, n) = max(dp(idx+1, n), sum[idx] + dp(idx+K, n-1))

 구간 arr[idx...idx+K-1]을 선택하는 경우와 선택하지 않는 경우 두 경우가 존재한다. 해당 구간을 선택하지 않는 경우, 소형 기관차를 사용하지 않았으므로 dp(idx+1, n)이 된다. 반대로 해당 구간을 선택하는 경우, idx+K-1번 객차까지 끌었으므로 이제 idx+K번 객차부터 고려해야 하고, 소형기관차 하나를 사용했으므로 sum[idx] + dp(idx+K, n-1)이 된다. 기저 사례는 n = 0 인 경우 혹은 idx > N-K 로 남은 구간이 없는 경우이다. 총 시간 복잡도는 O(N).

반응형

---

#include <iostream>
#include <string.h>

using namespace std;

int N, K;
int cache[50001][3];
int arr[50001];
int sum[50001];

int dp(int idx, int n) {
    if (n == 0) return 0;
    if (idx > N-K) return 0;
    
    int& ret = cache[idx][n];
    if (ret != -1) return ret;
    
    ret = max(dp(idx+1, n), sum[idx]+dp(idx+K, n-1));
    
    return ret;
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL); cout.tie(NULL);
    
    memset(cache, -1, sizeof(cache));
    
    cin >> N;
    for (int i=0; i<N; i++) cin >> arr[i];
    cin >> K;
    
    for (int i=0; i<K; i++) sum[0] += arr[i];
    for (int i=1; i<N; i++) sum[i] = sum[i-1] + arr[i-1+K] - arr[i-1];
    
    cout << dp(0, 3);

    return 0;
}