Traveler( Atcoder 196-E )

문제 : https://atcoder.jp/contests/abc197/tasks/abc197_e

E - Traveler

 

문제 파악하기

수직선 위에 있는 공을 ID 순서대로 가져오는데 걸리는 최소 시간을 구하는 문제입니다. 출발은 0에서 시작하며 모든 공을 가진 다음에는 출발점인 0으로 돌아와야 합니다. 공은 한 곳에 여러 개 존재할 수 없지만 ID가 동일한 공은 여러 개 있을 수 있으며, ID가 작은 공부터 가져와야 합니다. 공의 개수가 최대 200,000개이기 때문에 완전탐색으로는 문제를 해결할 수 없습니다.

 

문제 해결하기

우선, 우리는 공의 ID에 따라 모으는 순서가 결정된다는 걸 알고 있습니다. 그렇기 때문에 우선 ID를 기준으로 오름차순 정렬을 해보면 같은 ID의 공이 합쳐지게 됩니다. 이렇게 되면 공을 파악하는데 훨씬 수월해집니다. 

 

그 다음으로 우리가 신경써야 할 공을 찾아봅시다. 사실 우리는 같은 ID를 가진 공은 딱 2개만 신경쓰면 됩니다. 바로 가장 왼쪽에 위치한 공과 가장 오른쪽에 위치한 공입니다. 공이 얼마나 많든 상관없이 같은 ID의 공을 모두 모았을 때, 위치하는 곳은 바로 해당 ID 공의 가장 왼쪽 혹은 오른쪽 위치입니다. 그렇기에 우리는 ID마다 공 위치를 확인해보고 가장 왼쪽 혹은 오른쪽에 위치한 공이라면 따로 위치를 기록해두고, 나머지는 결괏값에 미리 떨어진 거리를 더해주면 됩니다.

 

ID마다 가장 왼쪽과 오른쪽 공을 찾았다면 마지막으로 할 일은 0에서 공을 가져오는 과정을 탐색하면 됩니다. 이 때, 우리는 상태를 나타내는 상태 함수를 사용할 수 있습니다. 상태 함수를 만들 때, 필요한 매개변수는 어떤 것이 있을까요? 우선, 현재 몇 번째 ID인지 확인할 변수가 1개 필요합니다. 그리고 현재 위치가 해당 ID의 왼쪽인지 오른쪽인지 표시할 변수도 필요합니다. 지금까지 나온 매개변수를 가지고 상태 함수를 만들면 다음과 같습니다.

sv( idx, f ) >> idx번째 ID이며, 현재 위치가 f인 상태(f=0: 왼쪽 / f=1: 오른쪽)

 

위와 같은 상태 함수를 만들었으면 이제 할 일은 이동할 수 있는 다음 상태를 찾는 일입니다. 사실 갈 수 있는 경우의 수는 (idx+1) 위치의 왼쪽으로 이동하거나 오른쪽으로 이동하는 경우로 총 2가지 있습니다. 다만, 신경써야 할 것은 만약 다음 상태에서 0번, 즉 왼쪽에 위치하고 싶다면 우리는 먼저 오른쪽 위치로 이동해야 합니다. 다음 상태에서 왼쪽에 위치한다는 건, 오른쪽으로 이동해서 왼쪽으로 이동하면서 모든 공을 가지고 있는 상태를 의미하기 때문입니다. 마찬가지로 다음 상태에 1번, 즉 오른쪽에 위치하고 싶다면 왼쪽으로 먼저 이동해야 합니다.

 

이렇게 상태 함수를 만든 다음, 다음 상태로 이동하면서 적절하게 종료 조건과 메모이제이션을 사용하면 문제를 해결할 수 있습니다. 문제를 해결하는 소스코드는 다음과 같습니다.

 

소스코드

#include <stdio.h>
#include <string.h>
#include <vector>
#include <utility>
#include <algorithm>
#define NMAX 200010
#define lint long long int
#define PAIR pair<int, int>
using namespace std;
 
int N, X, C;
vector< PAIR > inp, v;
 
int sz;
lint ret;
lint dp[NMAX][2];
 
// p와 (p+1) 사이의 거리 반환
lint len(int p, int f, int ff) {
    lint a, b;
 
    if(f == 0) a = v[p].first;
    else a = v[p].second;
 
    if(ff == 0) b = v[p+1].first;
    else b = v[p+1].second;
 
    return abs(a-b);
}
 
// dp[idx][f]: idx번째 위치이며, f번 위치인 경우(f는 최소/최대)
lint sv(int idx, int f) {
    if(dp[idx][f] >= 0) return dp[idx][f];
 
    if(idx == sz-1) return dp[idx][f] = 0;
    else {
        // 0번 위치: 1번 위치 > 0번 위치로 이동한 경우 >> sv(idx+1, 0) + len(idx, f, 1)
        // 1번 위치: 0번 위치 > 1번 위치로 이동한 경우 >> sv(idx+1, 1) + len(idx, f, 0)
        lint ret = min( sv(idx+1, 0)+len(idx, f, 1), sv(idx+1, 1)+len(idx, f, 0));
        return dp[idx][f] = ret;
    }
}
 
int main() {
    // input
    scanf("%d", &N);
    for(int i=0;i<N;i++) {
        scanf("%d %d", &X, &C);
        inp.push_back({C, X});
    }
 
    // Color 기준 정렬
    sort(inp.begin(), inp.end());
 
    // Color의 최솟값과 최댓값 찾기
    v.push_back({0, 0});
    for(int i=0;i<N;i++) {
        if(i==0 or inp[i-1].first != inp[i].first) v.push_back({inp[i].second, inp[i].second});
        else {
            v.back().second = inp[i].second;
 
            // 미리 더하기
            ret += inp[i].second - inp[i-1].second;
        }
    }
    // 종료 조건
    v.push_back({0, 0});
 
    // 동적 프로그래밍
    sz = v.size();
    memset(dp, -1, sizeof(dp));
    printf("%lld", ret  + sv(0, 0));
 
}

후기

생각을 하게 만드는 동적 프로그래밍 문제입니다. 점화식이 생각나면 깔끔하게 풀리는 기분 좋은 문제라 재미있게 풀었습니다. 동적 프로그래밍을 연습하는 사람에게 적절한 문제라고 생각합니다.