다리( BOJ 9006 )

문제 : https://www.acmicpc.net/problem/9006

 

문제 파악하기

왼쪽 집에서 오른쪽 집으로 가는 모든 거리의 합을 최소화하기 위해 다리를 놓아야 할 곳을 찾는 문제입니다. 다리의 길이는 항상 2이기 때문에 다리 길이에 대해서는 신경쓸 필요가 없으며, 어느 위치에 다리를 놓을 것인지만 생각하면 됩니다. 다만, N의 최대 범위가 10⁶이며, 입력되는 집의 위치가 -10⁷ ~ 10⁷이기 때문에 단순히 모든 경우를 탐색하는 방법으로는 절대 해결할 수 없습니다. 대신, 중복되는 계산을 효과적으로 줄여서 정답을 찾을 수 있습니다.

 

문제 해결하기

만약 우리가 h 위치에 다리를 놓았다고 가정하면, 우리가 최솟값을 구해야할 수식 f(h)는 다음과 같습니다. 참고로 수식에서 배열 a과 b는 각각 왼쪽과 오른쪽에 위치한 집의 위치를 의미합니다.

f(h) = M*( |a[1]-h| + ... +|a[N]-h| ) + N*( |b[1]-h| + ... + |b[M]-h| )

 

위의 수식에서 주목할 점은 절댓값 계산입니다. h의 값에 따라 수식은 (a[i]-h) 혹은 (h-a[i])로 바뀝니다. 하지만 같은 수식을 사용하는 값들은 한번에 구할 수 있다는 점에 주목해야 합니다. 만약 N=5이며, a[3]까지 h보다 작거나 같다면 절댓값을 계산하는 수식 g(h)는 다음과 같이 작성할 수 있습니다.

g(h) = (h-a[1]) + (h-a[2]) + (h-a[3]) + (a[4]-h) + (a[5]-h) = h*3 - (a[1]+a[2]+a[3]) + (a[4]+a[5]) - h*2

 

수식이 계산하기 간편하게 바뀐 걸 확인했나요? 절댓값 계산은 h의 크기에 따라 배열의 인덱스를 적절하게 설정하면 누적합으로 구할 수 있습니다. h보다 큰 숫자가 처음으로 나오는 위치를 확인한 다음, h보다 작은 값들의 누적합은 h에서 빼고 나머지 값들은 누적합에서 h를 빼면 됩니다. 이렇게 값을 계산하면 O(logN)만에 모든 수식값을 구할 수 있습니다.

 

그렇다면 정답이 될 수 있는 가능성이 있는 숫자들은 어떤 걸까요? 바로, 왼쪽과 오른쪽에 위치한 집의 위치입니다. 집에 위치가 아닌 곳에 다리를 만들게되면 집 위치에 다리를 만들었을 때의 최솟값보다 항상 결괏값이 커지게 됩니다. 간단하게 생각해보면, 더해야하는 값이 한 곳 늘기 때문이라고 할 수 있습니다. 따라서, 우리는 (N+M)개의 숫자 중에 정답이 있으니 (N+M)개의 숫자들만 비교하면 됩니다.

 

사실 좀 더 생각해보면 꼭 (N+M)개의 숫자를 모두 구할 필요는 없습니다. 위의 f(h)를 그래프로 그려보면 볼록한 함수가 그려집니다. 결국 우리는 (N+M)개의 숫자를 정렬한 다음, 변곡점을 지나는 순간 탐색을 종료하면 됩니다. 자세한 소스코드는 다음과 같습니다.

 

소스코드

#include <stdio.h>
#include <vector>
#include <algorithm>
#define NMAX 1000010
#define lint long long int
using namespace std;
 
int T;
lint t, ret;
lint inp[2][NMAX], pSum[2][NMAX];
vector< lint > num;
 
int idx;
int cnt[2];
unsigned lint tmp, vmin;
 
int main() {
    scanf("%d", &T);
    while(T--) {
        // init
        num.clear();
        
        // input
        scanf("%d %d", &cnt[0], &cnt[1]);
        for(int inpI=0;inpI<2;inpI++) {
            for(int j=1;j<=cnt[inpI];j++) {
                scanf("%lld", &inp[inpI][j]);
                
                num.push_back(inp[inpI][j]);
            }
        }
        
        // sort
        sort(num.begin(), num.end());
        for(int inpI=0;inpI<2;inpI++) {
            sort(inp[inpI]+1, inp[inpI]+cnt[inpI]+1);
            
            for(int j=1;j<=cnt[inpI];j++) pSum[inpI][j] = pSum[inpI][j-1] + inp[inpI][j];
        }
        
        // search
        vmin = 0xFFFFFFFFFFFFFFFF;
        for(int numI=0;numI<num.size();numI++) {
            
            tmp = 0;
            for(int inpI=0;inpI<2;inpI++) {
                auto tidx = upper_bound(inp[inpI]+1, inp[inpI]+cnt[inpI]+1, num[numI]);
                idx = (int)(tidx - inp[inpI]);
                
                
                // 건너편 집 개수만큼 곱해주기 >> cnt[1-inpI]
                // 현재 다리보다 아래쪽에 있는 모든 집과 다리까지 거리의 합
                tmp += cnt[1-inpI]*( num[numI]*(idx-1) - pSum[inpI][idx-1] );
                
                // 현재 다리보다 위쪽에 있는 모든 집과 다리까지 거리의 합
                tmp += cnt[1-inpI]*( (pSum[inpI][cnt[inpI]]-pSum[inpI][idx-1]) - num[numI]*(cnt[inpI]-idx+1) );
            }
            
            // 종료조건
            if(vmin < tmp) break;
            else if(vmin == tmp) continue;
            else {
                vmin = tmp;
                ret = num[numI];
            }
        }
        
        // print
        printf("%lld.0\n", ret);
    }
}
 

후기

자잘한 구현에서 실수를 해서 생각보다 오래 걸렸던 문제입니다. 특히 계산한 값의 범위가 long long int형보다 살짝 커서 unsinged long long int형을 사용해야 한다는 걸 늦게 떠올렸던 점이 발목을 잡았습니다. 어려운 자료구조 없이 최적화를 연습할 수 있어서 재미있던 문제입니다.