우체국

문제

수직선과 같은 일직선상에 N개의 마을이 위치해 있다. i번째 마을은 X[i]에 위치해 있으며, A[i]명의 사람이 살고 있다.

이 마을들을 위해서 우체국을 하나 세우려고 하는데, 그 위치를 어느 곳으로 할지를 현재 고민 중이다. 고민 끝에 나라에서는 각 사람들까지의 거리의 합이 최소가 되는 위치에 우체국을 세우기로 결정하였다. 우체국을 세울 위치를 구하는 프로그램을 작성하시오.

각 마을까지의 거리의 합이 아니라, 각 사람까지의 거리의 합임에 유의한다

입력

첫째 줄에 N(1 ≤ N ≤ 100,000)이 주어진다. 다음 N개의 줄에는 X[1], A[1], X[2], A[2], …, X[N], A[N]이 주어진다. 범위는 |X[i]| ≤ 1,000,000,000, 0 ≤ A[i] ≤ 1,000,000,000 이며 모든 입력은 정수이다.

출력

첫째 줄에 우체국의 위치를 출력한다. 가능한 경우가 여러 가지인 경우에는 더 작은 위치를 출력하도록 한다.

예제 입력 1

3
1 3
2 5
3 3

예제 출력 1

2

문제 풀이

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class BJ2141_Post {
    public static void main(String[] args) throws IOException {
        BufferedReader br=new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        long N=Integer.parseInt(br.readLine());
        PriorityQueue<Long[]> queue=new PriorityQueue<>(Comparator.comparing(o1->o1[0]));

        long total=0;

        for(long i=0;i<N;i++){
            st=new StringTokenizer(br.readLine());
            long a=Long.parseLong(st.nextToken());
            long b=Long.parseLong(st.nextToken());

            total+=b;
            Long[] l=new Long[2];
            l[0]=a;
            l[1]=b;
            queue.add(l);
        }

        //양팔 저울 생각하여 양 옆의 무게의 차이가 최소일 때 중심을 찾을 수 있음
        //이때, 총 인구수의 중간값쯤에서 무게차이가 최소일 것을 예상할 수 있기 때문에
        //총 인구수의 중간값을 지나면 break해줌
        long tmp=0;
        long ans=queue.peek()[0];
        for(long i=0;i<N;i++){
            Long[] l=queue.poll();
            tmp+=l[1];

            if(tmp>=(long) Math.ceil((double)total/(double)2)){
                ans=l[0];
                break;
            }
        }

        System.out.print(ans);
    }
}

이 문제를 이해하는데 시간이 꽤 걸렸고 풀이를 생각하는데에도 시간이 걸렸다.

우선, 문제를 분석하자면 다음으로 나타낼 수 있을 것 같다.

• 마을에 우체국을 세운다.
• 각 마을 거리 * 인구수의 값이 최소인 경우를 찾는다.

입력 제한조건을 보면 1,000,000,000까지의 범위이기 때문에 완전탐색을 통해 우체국의 위치를 찾게 되면 시간초과가 발생할 것이다.

시간초과를 피하기 위해 간단하게 답을 찾을 수 있는 방법을 찾아야했는데 며칠동안 고민한 덕분에 문제를 해결할 수 있었다.

예시를 들어 설명해보자면

4
1 5
2 5
3 5
4 5

위와 같은 경우를 생각하면 1 2 3 4의 위치에 모두 같은 인구수가 배치되어있다.

각 경우를 계산하면, 1*5+2*5+3*5=30 , 1*5+1*5+2*5=20, 1*5+1*5+2*5=20, 1*5+2*5+3*5=30 이 나오게 되고

계산값이 각각 30,20,20,30으로 대칭되어 나오게 된다는 사실을 파악할 수 있다. 그렇다면 우체국의 위치는 중간값인 2.5에 위치하게 될 것이라고 예상할 수 있는데, 마을에 위치해야하기 때문에 2에 위치하게 될 것이다.

이때, 우리는 양팔 저울의 무게중심을 떠올릴 수 있다.

무게중심을 기준으로 양옆의 무게가 같으면 수평을 만들 수 있게 되는데 우리는 인구수를 무게에 비유하여 문제를 해결해볼 수 있다.

또 다른 예시를 들어보자.

4
1 5
2 5
3 5
4 6

처음 예시에서 마지막 값만 인구수를 6으로 바꾼 경우이다.

문제의 방식처럼 값을 구하게 되면 33, 22, 21, 30의 값이 나오게 된다. 그렇다면 3의 위치에 우체국을 세울 수 있게되는데

처음 예시와 비교해보자면 오른쪽 끝에 인구수가 1 늘었는데 우체국의 위치 또한 오른쪽으로 이동했다.

양팔 저울의 무게중심을 떠올려보면 오른쪽에 무게가 늘어 무게중심이 오른쪽으로 움직인 것과 같다.

즉, 우체국의 위치를 중심으로 양 옆의 인구수를 구하여 차이가 최소인 경우를 구하면 그때가 우체국의 위치인것이다.

우체국의 양 옆의 인구수 구하기

long tmp=0;
long ans=queue.peek()[0];
for(long i=0;i<N;i++){
    Long[] l=queue.poll();
    tmp+=l[1];
    
    if(tmp>=(long) Math.ceil((double)total/(double)2)){
        ans=l[0];
        break;
    }
}

코드를 보면 의아할 수 있다. 양 옆의 인구수를 구하는데 왜 반복문이 하나밖에 없는지.

생각해보면 우체국의 위치를 모두 구해볼 필요가 없다.

양 옆의 인구수가 최소가 되려면 왼쪽 인구수가 전체 인구수의 절반즈음 되었을 때 양 옆의 인구수가 최소가 될 것임으로 예상할 수 있다.

코드를 보면 tmp 변수에 왼쪽 인구수의 값을 더해주고 ans에는 우체국의 위치값을 넣어준다.

왼쪽부터 하나씩 들어오면서 왼쪽 인구수가 전체 인구수의 절반이상이 되는 경우에 탐색을 그만둔다.

ans에는 인덱스 i 값이 아닌 배열에 담긴 마을의 위치 l[0] 값을 넣어줘야한다!

왜 절반이상이 되었을때 멈추어야하는지는 다음 예시를 생각해 볼 수 있다.

2
1 1
2 2

위 경우의 답은 2가 되어야한다. 왼쪽 값이 절반 이하인 경우를 찾게되면 답은 1이라고 도출하게 되기 때문에 정확한 위치를 찾기 위해서는 절반 이상일때 break를 해주어야한다.

우선순위 큐

PriorityQueue<Long[]> queue=new PriorityQueue<>(Comparator.comparing(o1->o1[0]));

그리고 각 마을의 위치와 인구수를 넣어주기 위해 우선순위큐를 사용했다. 마을 위치가 순서대로 들어온다는 보장이 없기 때문에 정렬을 해주어야했는데 우선순위큐를 이용하여 시간복잡도를 줄였다.

이때, 배열값을 넣어주기 때문에 마을의 위치를 기준으로 오름차순으로 정렬되도록 Comparator.comparing(o1->o1[0])을 넣어줬다.

이런 과정을 거치면 답을 도출할 수 있다.

처음에 인구수의 절반을 지나친다는 생각을 하지 못해 일일이 양 옆의 인구수를 모두 구해줬더니 당연히 시간초과가 발생했다.

시간초과 발생한 코드

public static void main(String[] args) throws IOException {
        BufferedReader br=new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        int N=Integer.parseInt(br.readLine());
        long[][] loc=new long[N][2];

        for(int i=0;i<N;i++){
            st=new StringTokenizer(br.readLine());
            loc[i][0]=Long.parseLong(st.nextToken());
            loc[i][1]=Long.parseLong(st.nextToken());
        }

        //마을위치 정렬
        Arrays.stream(loc).sorted();
        long ans=Integer.MAX_VALUE;
        long idx=0;

        for(int i=0;i<N;i++){
            long tmp_l=0;
            long tmp_r=0;
            for(int j=0;j<i;j++){
                tmp_l+=loc[j][1];
            }

            for(int k=i+1;k<N;k++){
                tmp_r+=loc[k][1];
            }

            if(Math.abs(tmp_r-tmp_l)<ans){
                ans=Math.abs(tmp_r-tmp_l);
                idx=i;
            }
        }

        System.out.println(idx+1);

    }

long 배열로 마을과 인구수의 정보를 받고 마을 위치를 기준으로 정렬해준 다음

양 옆의 인구수를 모두 구하여 최소값을 일일이 구하는 방식을 썼다. 당연하게도 시간초과가 발생했던 코드이다.