소형 기관차 (Gold 4)
문제
2616번: 소형기관차
첫째 줄에 기관차가 끌고 가던 객차의 수가 입력된다. 그 수는 50,000 이하이다. 둘째 줄에는 기관차가 끌고 가던 객차에 타고 있는 손님의 수가 1번 객차부터 차례로 입력된다. 한 객차에 타고 있
www.acmicpc.net
접근법
문제를 보고 처음 접근할 때는 완전 탐색이 가능한지를 먼저 살펴보는 것이 좋을 때가 많습니다. 완전 탐색으로도 쉽게 풀릴 수 있는 문제를 어렵게 고민하는 경우가 생길 수도 있기 때문입니다. 그리고 완전 탐색으로 답을 구할 수 없더라도, 완전 탐색으로 접근한 방법을 최적화하여 문제를 해결하는 경우도 많습니다. 이번 문제도 완전 탐색으로 먼저 접근해 보았습니다.
소형 기차는 총 3대입니다. 기차의 전체 객실 수가 N개, 소형기관차의 길이가 L이라고 했을 때 첫 번째 소형 기관차가 선택할 수 있는 열차칸의 경우의 수는 N - L입니다. 두 번째 세 번 째는 각각 N - 2L, N - 3L 개의 경우의 수를 가집니다. 따라서 모든 경우의 수를 계산하기 위해서는 최악의 경우 \(O(N^3)\)의 계산이 필요하며 N의 최댓값이 50,000으로 상당히 크기 때문에 이번 문제는 완전 탐색으로 풀 수 없는 문제입니다.
계산 시간을 줄이기 위해서 이번 문제는 동적계획법으로 중복된 계산을 줄여줄 수 있습니다. 먼저 소형 기관차가 한 대인 경우를 생각해봅시다. 소형 기관차가 한대인 경우에는 전체 객실 중 연속된 L개의 객실 중 가장 최댓값을 선택하면 됩니다. N칸의 객실들이 있을 때 0부터 임의의 i번째 칸 까지 최댓값을 dp[0] 배열에 다음과 같이 계산하여 저장할 수 있습니다.
// 소형기관차가 한대일 경우
for(int i = L - 1; i < N; ++i)
{
dp[0][i] = max(dp[0][i - 1], sum[i]);
}
소형 기관차가 두대일 경우를 생각해 봅시다. 소형 기관차가 두대일 경우에는 두 가지 경우를 따져봐야 합니다. 두 소형기관차 중 한 대가 i 번째 칸 객실을 선택하는 경우와 i 번째 칸 객실을 선택하지 않는 경우로 나눠질 수 있습니다. 만약 i번째 칸 객실 선택할 경우 다른 소형 기관차는 한대는 0 ~ i-L 중 가장 최댓값을 선택하는 것이 효율적일 것입니다. 그리고 이 값은 앞에서 계산한 dp[0][i-L] 값입니다. i 번째 칸 객실을 선택하지 않는다면, i-1 번째 칸 객실을 선택한 경우와 동일해집니다. 따라서 소형 기관차가 두 대일 경우는 다음과 같이 계산할 수 있습니다.
// 소형기관차가 두대일 경우
for (int i = (2 * L) - 1; i < N; ++i)
{
dp[1][i] = max(dp[1][i - 1], sum[i] + dp[0][i - L]);
}
이제 소형 기관차가 세대일 경우를 생각해 봅시다. 소형 기관차가 세대일 경우에도 두 가지 경우를 따져봐야 합니다. 소형 기관차 중 한 대가 i 번째 칸 개실을 선택하는 경우와 i 번째 칸 개실을 선택하지 않는 경우로 나눠질 수 있습니다. 만약 i번째 칸 객실을 선택할 경우 다른 소형 기관차 두대는 0 ~ i - *L 중 가장 최댓값을 선택하는 것이 효율적일 것입니다. 그리고 이 값은 앞에서 계산한 dp[1][i-L] 값입니다. i 번째 칸 객실을 선택하지 않는다면, i-1 번째 칸 객실을 선택한 경우와 동일해집니다. 따라서 소형 기관차가 세 대일 경우는 다음과 같이 계산할 수 있습니다.
// 소형기관차가 세대일 경우
for (int i = (3 * L) - 1; i < N; ++i)
{
dp[2][i] = max(dp[2][i - 1], sum[i] + dp[1][i - L]);
}
그리고 dp[2][N-1]가 구해야 하는 값이 됩니다. 정리하면 아래와 같이 계산됩니다. 입력값은 문제에서 주어진 예시에서 40짜리 객실이 하나 더 추가하였습니다.
in 35 40 50 10 30 45 60 40
sum 0 75 90 60 40 75 105 100
dp0 0 75 90 90 90 90 105 100
dp1 0 0 0 135 135 165 195 195
dp2 0 0 0 0 0 210 240 265
이렇게 시간 복잡도 \(O(N)\), 공간 복잡도 \(O(N)\)으로 계산을 완료할 수 있었습니다.
전체 코드
#include <iostream>
#include <algorithm>
#include <array>
#include <queue>
#define endl '\n'
using namespace std;
constexpr int MAX_LEN = 50'000;
int main()
{
//입출력 성능향상을 위한 설정
ios_base::sync_with_stdio(false);
cout.tie(NULL);
cin.tie(NULL);
int N; // N: 50,000 이하
int L; // 소형기관차의 길이
array<int, MAX_LEN> arr{};
array<array<int, MAX_LEN>, 3> dp{};
array<int, MAX_LEN> sum{};
// Input Data
cin >> N;
for (int i = 0; i < N; ++i)
cin >> arr[i];
cin >> L;
// 소형기관차가 i번 열차까지 끌때 인원수 합 계산
queue<int> q;
int temp{};
for (int i = 0; i < N; ++i)
{
temp += arr[i];
q.push(arr[i]);
if (i < L - 1) continue;
sum[i] = temp;
temp -= q.front();
q.pop();
}
// 소형기관차가 한대일 경우
for(int i = L - 1; i < N; ++i)
{
dp[0][i] = max(dp[0][i - 1], sum[i]);
}
// 소형기관차가 두대일 경우
for (int i = (2 * L) - 1; i < N; ++i)
{
dp[1][i] = max(dp[1][i - 1], sum[i] + dp[0][i - L]);
}
// 소형기관차가 세대일 경우
for (int i = (3 * L) - 1; i < N; ++i)
{
dp[2][i] = max(dp[2][i - 1], sum[i] + dp[1][i - L]);
}
cout << dp[2][N-1];
return 0;
}