[Baekjoon][Python] 18310번 안테나

18310번 안테나

문제

풀이

문제를 보자마자 생각난 것은, 이진탐색을 활용한 풀이였다.
먼저 중앙값을 기준으로 좌, 우로의 합을 구해 더 작은 쪽으로 움직이는 것이다. 완전 이진 탐색보다는 시간복잡도가 높지만 그나마 괜찮은 방법으로 생각했던 것이었다.
하지만, 문제의 주요한 조건을 하나 놓쳐서 너무 복잡하게 푼 것 같다. 바로 안테나는 집이 위치한 곳에만 위치할 수 있다는 것이다.

import sys

def totalSum(loc, locations):
    totalSum = 0
    for location in locations:
        totalSum += abs(loc-location)
    return totalSum
def findPoleLocation(locations, n):
    poleLocation = locations[n // 2]
    if poleLocation - 1 <= 0 or poleLocation + 1 > locations[-1]:
        return poleLocation
    startSum = totalSum(poleLocation, locations)
    left = totalSum(poleLocation - 1, locations)
    right = totalSum(poleLocation + 1, locations)
    flag = 0
    minSum = 0
    if startSum >= left:
        flag = -1
        poleLocation -= 1
        minSum = left
    elif startSum <= right:
        return poleLocation
    elif startSum > right:
        flag = 1
        poleLocation += 1
        minSum = right

    while True:
        sumData = totalSum(poleLocation + flag, locations)
        if minSum >= sumData:
            if flag == 1:
                return poleLocation
            minSum = sumData
            poleLocation += flag
        else:
            return poleLocation

# 집의 수
N = int(sys.stdin.readline())

locations = list(map(int, sys.stdin.readline().split()))
locations.sort()

위 코드에서 집이 위치한 곳에만 안테나를 위치시킨다면 조금 더 시간을 단축시킬 수 있지만 더 좋은 생각이 떠올랐다.

바로 정렬 후에 가운데 점을 선택하는 것이다.
가운데 값을 선택하게 될 때, 거리의 합이 최소가 된다는 것이다.

따라서 아래와 같이 매우 간단한 코드로 가능해진다.

import sys
N = int(sys.stdin.readline())
locations = list(map(int, sys.stdin.readline().split()))
locations.sort()
print(locations[(N - 1) // 2])

결론

이번 문제에서 느낀 것은, 문제를 더 자세하게 읽고 조금 더 창의적으로 생각해보자이다. 정렬이 핵심인 문제였다.