백준 18310 안테나

문제

일직선 상의 마을에 여러 채의 집이 위치해 있다. 이중에서 특정 위치의 집에 특별히 한 개의 안테나를 설치하기로 결정했다. 효율성을 위해 안테나로부터 모든 집까지의 거리의 총 합이 최소가 되도록 설치하려고 한다. 이 때 안테나는 집이 위치한 곳에만 설치할 수 있고, 논리적으로 동일한 위치에 여러 개의 집이 존재하는 것이 가능하다.

집들의 위치 값이 주어질 때, 안테나를 설치할 위치를 선택하는 프로그램을 작성하시오.

예를 들어 N=4이고, 각 위치가 1, 5, 7, 9일 때를 가정하자.

이 경우 5의 위치에 설치했을 때, 안테나로부터 모든 집까지의 거리의 총 합이 (4+0+2+4)=10으로, 최소가 된다.

---

코드

import java.io.*;
import java.util.*;
import java.util.stream.Collectors;

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        //N 입력 받기
        int n = Integer.parseInt(br.readLine());

        //Map에 (X좌표, 갯수)로 집 저장
        Map<Integer, Integer> houses = new HashMap<>();
        StringTokenizer st = new StringTokenizer(br.readLine());
        for(int i=0; i<n; i++){
            int next = Integer.parseInt(st.nextToken());
            houses.put(next, houses.getOrDefault(next, 0)+1);
        }

        //중복이 제거된 X좌표 오름차순 정렬
        List<Integer> list=houses.keySet().stream().sorted().collect(Collectors.toList());

        //첫 번째 집에 있을 때의 총 거리 sum, 오른쪽에 있는 집들의 갯수 rights
        long sum=0;
        for(int i:list)
            sum+=(long)houses.get(i)*(i-list.get(0));
        int rights=n-houses.get(list.get(0));

        //가장 작은 총 거리를 갱신하기 위한 minSum, 해당 집의 인덱스인 answer
        long minSum=sum;
        int answer=list.get(0);

        //두 번째 집부터 끝까지
        for(int i=1; i<list.size(); i++){
            //이전 집과의 거리 차이
            int difference=list.get(i)-list.get(i-1);

            //(현재 집보다 오른쪽에 있는 집들의 갯수*거리 차이)만큼 sum에서 빼주고, 왼쪽의 집만큼 더해준다
            sum+=(long)(n-2*rights)*difference;
            
            //오른쪽 집들의 갯수 갱신
            rights-=houses.get(list.get(i));

            //정답 갱신
            if(minSum>sum){
                minSum=sum;
                answer=list.get(i);
            }
        }

        System.out.println(answer);
    }
}
//sum-=rights*difference
//sum+=(n-rights)*difference

//2*sum-2*rights*A+n*A
//A=next-list.get(i-1)

---

해결 과정

1. 각 집의 위치에서 모두 거리를 계산하면 O(N^2)이 걸리므로 최대 20만개의 집은 시간초과가 뜬다. 오름차순 정렬 후, 한 집을 기준으로 잡고 다음 집으로 옮겨가는 거리 차이만큼만 갱신해주면 O(N)으로 풀 수 있다.

2. 😁