문제

1보다 큰 N개의 도시 중 한 곳에 공항을 지을 예정입니다. 사람들의 편의를 위해 공항으로부터 각 사람들까지의 도시간 이동 거리가 최소가 되는 도시에 짓기로 하였습니다. 편의상 도시는 일직선상에 놓여있다고 가정하며 좌표의 범위는 음수가 포함됩니다. 또한 좌표는 정렬되어 있지 않습니다. 직선상의 위치와 그 도시에 사는 사람들의 수가 주어질 때, 공항을 지을 도시의 위치를 반환해주는 함수 chooseCity 함수를 완성하세요. 거리가 같은 도시가 2개 이상일 경우 위치가 더 작은 쪽의 도시를 선택하면 됩니다. 예를 들어 다음과 같은 정보의 도시가 있다고 가정해 봅시다.

위치 1 2 3
인구수 5 2 3

이 살 경우, 각각의 도시에 공항을 지었을 때의 사람들의 이동 거리는 8, 8, 12 이므로 1번 또는 2번에 지을 수 있지만, 1의 위치가 더 작으므로 1을 반환해주면 됩니다.

풀이 1

import java.util.*;

class TryHelloWorld
{
    public int chooseCity(int n, int [][]city)
    {
        int answer = 0;
        int sum = 0;
        int tmp = 0;

        for(int i=0; i<n; ++i){
            for(int j=0; j<n; ++j){
            	tmp += Math.abs(city[i][0]-city[j][0])*city[j][1];
            }

            if(tmp < sum || i == 0){
                sum = tmp;
                answer = city[i][0];
            }else if(tmp == sum && answer > city[i][0]){
            	answer = city[i][0];
            }

            tmp = 0;
        }

        return answer;
    }
  
  public static void main(String[] args)
  {
    TryHelloWorld test = new TryHelloWorld();
    int tn = 3;
    int [][]tcity = {{1,5},{2,2},{3,3}};
    System.out.println(test.chooseCity(tn,tcity));
  }
}

제일 처음에 풀었던 방식이다. O(N^2^)의 시간복잡도를 가지게 된다. 간단하게 풀이해보자면, 각 위치별로 도시간의 거리를 구한 후(이중 for문), 각 도시별로 구한 거리끼리 비교하여(if문) 거리가 작은 쪽의 위치를 answer에 넣도록 하였다.

이 풀이로 예제는 통과하지만 제출을 하게되면 수많은 테스트 케이스에서 타임아웃에 걸리게 된다. 계속 풀어보다가 좌절하고 결국 구글의 힘을 빌려보기로 하였다.ㅠㅠ

풀이 2

import java.util.*;

class TryHelloWorld
{
  public int chooseCity(int n, int [][]city)
  {
    int answer = city[0][0];
    int left = 0;
    int right = 0;

    Arrays.sort(city, new Comparator<int[]>() {
      public int compare(int[] arr1, int[] arr2) {
      	return Integer.compare(arr1[0], arr2[0]);
      }
    });

    for(int i=0; i<n; ++i) right += city[i][1];

    for(int i=0; i<n-1; ++i){
      left += city[i][1];
      right -= city[i][1];
      if(right > left) answer = city[i+1][0];
    }

    return answer;
  }
  
  public static void main(String[] args)
  {
    TryHelloWorld test = new TryHelloWorld();
    int tn = 3;
    int [][]tcity = {{1,5},{2,2},{3,3}};
    System.out.println(test.chooseCity(tn,tcity));
  }
}

구글링을 해보아도 자바 풀이 중에는 제대로 된 것을 찾을 수 없었다. 또한, 어떻게든 제출이 되는 풀이를 내고 다른 사람의 풀이를 보았는데 분명 통과된 소스인데 내가 다시 돌려보면 타임아웃이 나거나 정답이 아니라고 뜨더라… 아마 운이 좋게 테스트 케이스를 통과하여 제출이 된 모양이다.

그래서 자바 외의 파이썬 언어로 푼 풀이를 찾아보았다. 그 중에 하나의 풀이이다. 먼저 위치를 기준으로 오름차순 정렬을 하였다.(Arrays.sort) 이유는 중간 위치가 기준이 되기 때문이다. 그 다음에 가장 왼쪽 도시에 공항을 지었을 겯우를 가정하여 left엔 0, right엔 모든 인구의 합을 할당하였다.

각 인덱스에 해당하는 인구수만큼 left는 더하고(좌측의 인구수 증가) right(좌측만큼 인구수 감소)는 뺀다. 이유를 찾아보니 인구수가 평균에 가까워질수록, 혹은 클수록 최적의 이동 거리를 만들어내기 때문이다.

이대로 풀게 되면 n log n의 시간 복잡도를 가지게 되어 타임아웃도 나지 않고 정상적으로 정답 제출이 된다.