[KOI] 백준 2616 - 소형기관차

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

 

다음 dp 테이블을 생각한다.

 

dp[i][j] := j번째 객차까지 소형 기관차 i개를 썼을 때 수송할 수 있는 최댓값

 

이렇게 했을 때 지금 객차부터 소형 기관차를 하나 사용하여 사람들을 수송했을 때, 지금 객차에서 소형 기관차를 사용하지 않고 다음 객차로 넘어갔을 때를 메모이제이션하여 문제를 풀 수 있다.

 

일일이 객차의 인원을 더하는 것은 무리이므로 객차별 인원을 누적합으로 관리하면 제한 시간 내에 정답을 받을 수 있다.

 

전체 코드

#pragma GCC target("avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

#include <bits/stdc++.h>

using namespace std;

int n;
int ar[50001];
int prefix[50001];
int dp[4][50001];
int lmt;

int memo(int train, int idx)
{
  if (idx > n) return -1e9;

  if (idx == n)
  {
    if (train != 3) return -1e9;
    return 0;
  }

  int &ret = dp[train][idx];

  if (ret != -1) return ret;
  ret = 0;
  int tmp = memo(train, idx + 1);

  if (train < 3)
  {
    int nxt = lmt + idx;

    if (nxt <= n)
    {
      int val = prefix[nxt] - prefix[idx];
      tmp = max(tmp, val + memo(train + 1, nxt));
    }
  }

  return ret += tmp;
}

void solve()
{
  memset(dp, -1, sizeof dp);
  cin >> n;
  for (int i = 1; i <= n; ++i)
  {
    cin >> ar[i];
    prefix[i] = ar[i] + prefix[i - 1];
  }

  cin >> lmt;

  cout << memo(0, 0);
}

int main()
{
  cin.tie(nullptr);
  ios::sync_with_stdio(false);

  solve();
}