[Baekjoon] 가운데를 말해요(python)

https://www.acmicpc.net/problem/1655

문제 접근

이 문제를 접근할 때 가장 단순하게 생각하면 리스트에 숫자를 넣으면서 Sorting 후 중간값을 출력하면 된다고 생각할 것이다. 하지만 제한 시간은 0.1초이고 파이썬 정렬의 시간 복잡도는 O(NlogN)이므로 시간 초과가 나올 것이다.

또 생각했던 방법은 중간값을 저장하는 변수 Center를 두고 Left, Right라는 리스트를 통해 Center보다 작으면 Left에 크면 Right에 넣는 방법을 생각했었다. 하지만 그렇게 되면 중간값을 갱신할 방법이 없었고 다른 사람의 풀이를 통해 Left와 Right를 단순 리스트가 아닌 Heap을 이용하면 중간값을 갱신할 수 있는 솔루션이 될 수 있었다. 따라서 Left를 Max Heap, Right를 Min Heap으로 설정해서 문제를 풀었다.

 

---

풀이 과정

위에서 언급했던 것처럼, 이 문제를 풀기 위해서 중간값을 저장하는 변수를 두고 중간값보다 작으면 Left 크면 Right에 넣어주면 된다.

Left는 Max Heap, Right는 Min Heap으로 설정해서 문제를 해결하면 되는데 중간값을 구하려면 Left의 최댓값과 Right의 최솟값을 비교하면 중간값을 알 수 있기 때문에 Center라는 변수를 따로 저장할 필요없이 두 개의 Heap만을 이용하면 Left의 최댓값을 출력해주면 된다.

 

숫자를 하나씩 Heap에 넣게 될텐데 Left에 넣을지 Right에 넣을지 결정할 수 있는 로직이 필요하다. 숫자를 넣을 때 우선적으로 Left에 넣고 두 개의 Heap 크기가 1보다 더 크지 않도록 해야 한다.

또 Left의 최댓값이 Right의 최솟값보다 크면 두 숫자의 위치를 바꿔야 한다. 그 이유는 현재 Center라는 변수는 없지만 가상의 Center가 존재하는 것이기 때문에 중간값보다 작으면 Left 크면 Right에 넣어야 하기 때문이다.

---

정답 코드

import heapq
import sys

n = int(sys.stdin.readline())

left, right = [], []

for _ in range(n):
    number = int(sys.stdin.readline())

    if len(left) == len(right):
        heapq.heappush(left, -number)
    else:
        heapq.heappush(right, number)

    if right and -left[0] > right[0]:
        # 왼쪽에서 가장 큰 값이 오른쪽에서 가장 작은 값보다 크다면 교체
        min_val = heapq.heappop(right)
        max_val = heapq.heappop(left)

        heapq.heappush(left, -min_val)
        heapq.heappush(right, -max_val)

    print(-left[0])