BOJ 2616 - 소형기관차

문제 링크

문제#

기차는 맨 앞에 있는 기관차 1대가 손님이 탄 객차 여러 칸을 끌고 간다. 기관차가 고장나면 기차를 운행할 수 없게 되므로 최근 철도청은 기관차 고장에 대비하여 몇몇 역에 소형 기관차 3대를 배치하기로 결정하였다. 소형 기관차는 평소에 이용하는 기관차보다 훨씬 적은 수의 객차만을 끌 수 있다.

기관차가 고장났을 때 끌고 가던 객차 모두를 소형 기관차 3대가 나누어 끌 수 없기 때문에, 소형 기관차들이 어떤 객차들을 끌고 가는 것이 좋을까하는 문제를 고민하다가 다음과 같이 하기로 결정하였다.

1. 소형 기관차가 최대로 끌 수 있는 객차의 수를 미리 정해 놓고, 그보다 많은 수의 객차를 절대로 끌게 하지 않는다. 3대의 소형 기관차가 최대로 끌 수 있는 객차의 수는 서로 같다.
2. 소형 기관차 3대를 이용하여 최대한 많은 손님을 목적지까지 운송하도록 한다. 각 객차 마다 타고 있는 손님의 수는 미리 알고 있고, 다른 객차로 손님들이 이동하는 것은 허용하지 않는다.
3. 각 소형 기관차는 번호가 연속적으로 이어진 객차를 끌게 한다. 객차는 기관차 바로 뒤에 있는 객차부터 시작하여 1번 부터 차례로 번호가 붙어있다.

예를 들어 기관차가 끌고 가던 객차가 7칸이고, 소형 기관차 1대가 최대로 끌 수 있는 객차 수는 2칸이라고 하자. 그리고 1번 부터 7번까지 각 객차에 타고 있는 손님의 수가 아래 표와 같다고 하자. 괄호속에 있는 숫자는 객차 번호를 나타낸다.

소형 기관차 3대는 각각 1-2번, 3-4번, 그리고 6-7번 객차를 끌고 가면 손님 240명을 운송할 수 있고, 이보다 많은 수의 손님을 운송할 수 없다.

기관차가 끌고 가던 객차의 수와 각 객차에 타고 있던 손님의 수, 그리고 소형 기관차가 최대로 끌수 있는 객차의 수가 주어질 때, 소형 기관차 3대를 이용하여 최대로 운송할 수 있는 손님 수를 구하는 프로그램을 작성하시오.

입력#

첫째 줄에 기관차가 끌고 가던 객차의 수가 입력된다. 그 수는 50,000 이하이다. 둘째 줄에는 기관차가 끌고 가던 객차에 타고 있는 손님의 수가 1번 객차부터 차례로 입력된다. 한 객차에 타고 있는 손님의 수는 100명 이하이고, 입력되는 숫자들 사이에 빈칸이 하나씩 있다. 셋째 줄에는 소형 기관차가 최대로 끌 수 있는 객차의 수가 입력된다. 그 수는 기관차가 끌고 가던 객차 수의 1/3보다 적다.

출력#

한 줄에 소형 기관차 3대를 이용하여 최대로 운송할 수 있는 손님 수를 출력한다.

풀이#

연속 구간 합이 반복해서 필요하므로 누적합을 먼저 만든다. 이후 solve(train, car)를 "car번 객차부터 볼 때, train번째 기관차부터 얻을 수 있는 최대 승객 수"로 두는 DP를 사용한다.

각 위치에서는 두 가지 선택만 있다. 현재 객차를 건너뛰고 다음 칸으로 가거나, 지금부터 길이 k의 구간을 하나 선택하고 다음 기관차로 넘어가는 것이다. 두 경우 중 큰 값을 택하면 된다.

기관차는 총 3대뿐이라 상태 수가 크지 않고, 메모이제이션으로 충분히 해결된다.

현재 칸에서 기관차 하나를 쓰기로 했다면, 길이 k의 객차를 반드시 연속으로 묶어야 한다. 이 구간의 승객 수는 누적합으로 $O(1)$에 계산할 수 있고, 다음 상태는 train + 1, car + k가 된다. 반대로 지금 칸을 사용하지 않으면 car + 1로만 넘어가면 된다.

즉 이 문제는 "지금 여기서 한 대를 쓰는 것이 이득인가, 다음 칸으로 넘기는 것이 이득인가"를 반복해서 판단하는 구간 DP다. 기관차 수가 3대로 고정되어 있다는 점 덕분에 완전탐색처럼 보이는 경우의 수도 충분히 줄어든다.

코드#

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
 
int n, k;
vector<int> sum;
vector<vector<int>> dp;
 
int solve(int train, int car) {
	if (train > 3) return 0;
	if (car + k - 1 > n) return 0;
	if (dp[train][car] != -1) return dp[train][car];
 
	// 객차를 건너뛴 경우
	int skip = solve(train, car + 1);
 
	// 객차를 선택한 경우
	int select = (sum[car + k - 1] - sum[car - 1]) + solve(train + 1, car + k);
 
	return dp[train][car] = max(skip, select);
}
 
int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
	
	cin >> n;
	sum.resize(n + 1);
	dp.resize(4, vector<int>(n + 1, -1));
 
	int input;
	for (int i = 1; i <= n; i++) {
		cin >> input;
		sum[i] = sum[i - 1] + input;
	}
	cin >> k;
	
	cout << solve(1, 1) << '\n';
	return 0;
}

복잡도#

• 시간 복잡도: $O(3n)$
• 공간 복잡도: $O(3n)$

마무리#

기관차가 3대로 고정되어 있어서 상태 수가 작다. 누적합으로 구간 점수를 즉시 뽑고, 고를지 넘길지 DP로 정리하면 된다.