취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 2616번 문제인 소형기관차이다.
문제는 아래 링크를 확인하자.

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

dp1[x]를 1번 객차부터 x번 객차까지의 범위에서 K개의 연속한 객차를 하나의 소형기관차를 이용해 끌고갈 때 최대로 옮길 수 있는 손님의 수로 정의하자(단, x는 K보다 크다.). 소형기관차가 객차를 끄는 방법은 x번 객차를 끌거나 끌지 않는 두 가지로 구분할 수 있다. 각 경우의 옮길 수 있는 최대 손님의 수는 dp[x-1]과 dp[x-K], 그리고 x-K+1~x번 객차의 손님의 수의 합을 이용하면 계산해낼 수 있음을 관찰하자. 후자의 값은 prefix sum 배열을 미리 전처리해두는 것으로 빠르게 계산할 수 있다.

위와 같은 방식으로 dp2[x] 및 dp3[x] 또한 정의해 그 값들을 구해내자. 이 때 문제의 답은 dp3[N]이 될 것이다.

이 방법의 시간복잡도는 \(O(N)\)이다.

아래는 제출한 소스코드이다.

#include <iostream>
using namespace std;

int N, K;
int psum[50001];
int dp1[50001];
int dp2[50001];
int dp3[50001];

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

	cin >> N;
	for (int i = 1; i <= N; i++) {
		cin >> psum[i];
		psum[i] += psum[i - 1];
	}
	cin >> K;

	for (int i = K; i <= N; i++) dp1[i] = max(dp1[i - 1], psum[i] - psum[i - K]);
	for (int i = K * 2; i <= N; i++) dp2[i] = max(dp2[i - 1], dp1[i - K] + psum[i] - psum[i - K]);
	for (int i = K * 3; i <= N; i++) dp3[i] = max(dp3[i - 1], dp2[i - K] + psum[i] - psum[i - K]);

	cout << dp3[N];
}