[Python/Gold2/1655] 가운데를 말해요

|    문제

백준이는 동생에게 "가운데를 말해요" 게임을 가르쳐주고 있다. 백준이가 정수를 하나씩 외칠때마다 동생은 지금까지 백준이가 말한 수 중에서 중간값을 말해야 한다. 만약, 그동안 백준이가 외친 수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다.

예를 들어 백준이가 동생에게 1, 5, 2, 10, -99, 7, 5를 순서대로 외쳤다고 하면, 동생은 1, 1, 2, 2, 2, 2, 5를 차례대로 말해야 한다. 백준이가 외치는 수가 주어졌을 때, 동생이 말해야 하는 수를 구하는 프로그램을 작성하시오.

|    입력

첫째 줄에는 백준이가 외치는 정수의 개수 N이 주어진다. N은 1보다 크거나 같고, 100,000보다 작거나 같은 자연수이다. 그 다음 N줄에 걸쳐서 백준이가 외치는 정수가 차례대로 주어진다. 정수는 -10,000보다 크거나 같고, 10,000보다 작거나 같다.

|    출력

한 줄에 하나씩 N줄에 걸쳐 백준이의 동생이 말해야 하는 수를 순서대로 출력한다.

 

 

---

정답 코드

import sys, heapq
input = sys.stdin.readline

class Solution:
    def say_mid(self):
        N = int(input())
        arr = [int(input()) for _ in range(N)]
        res = [arr[0]]
        left, right = [], []
        mid = arr[0]
        for idx, num in enumerate(arr[1:], 1):
            if num < mid:
                heapq.heappush(left, -num)
            else:
                heapq.heappush(right, num)

            if idx % 2 == 1:
                if len(left) > len(right):
                    heapq.heappush(right, mid)
                    mid = -heapq.heappop(left)

            elif idx % 2 == 0:
                if len(left) > len(right):
                    heapq.heappush(right, mid)
                    mid = -heapq.heappop(left)
                elif len(left) < len(right):
                    heapq.heappush(left, -mid)
                    mid = heapq.heappop(right)
            res.append(mid)

        for r in res:
            print(r)

if __name__ == "__main__":
    s = Solution()
    s.say_mid()

 

 

|    비고

중앙값 말하기 문제의 하위 버전이라고 생각하면 된다. 중앙값보다 작은 값을 담을 left heapq, 중앙값보다 큰 값을 담을 right heapq를 사용해서 풀어야 한다. 중앙값을 비교해야하기 때문에, 작은 값(left)와 큰 값(right)의 값을 현재 중앙값과 비교해서 다음 중앙값을 정하는 방식이다. 이때 작은 값을 담는 heapq는 최대힙, 큰 값을 담는 heapq는 최소힙이어야 하는데 왜일까?

예를 들어 중앙값은 5이고 left엔 [1,2,3], right엔 [10]이 담겨져 있다고 치자. 그럼 다음 중앙값은 무엇이 될까? 당연히 3이 될 것이다. 그럼 3을 mid(중앙값)으로 지정하려면 left에서 3을 꺼내야 하는데, 이때 최대힙을 사용해서 꺼내면 된다.

또한, 중앙값이 3이고 left엔 [1,2], right엔 [4,5,10]이 담겨져 있다면, right에 있는 4가 다음 중앙값이 될 것이다. 이때는 최소힙을 사용하여 4를 꺼낼 수 있다.