가운데를 말해요

숫자들이 차례로 쌓여감에 따라 숫자들의 중간값을 출력하는 문제다.

리스트를 만들어 숫자들을 집어넣고 매번 정렬을 해 중간값을 찾게 되면 시간 복잡도가 급격히 증가한다. (이 문제에서 주어지는 정수의 최대 개수는 100,000개 이다.)

따라서 heap을 사용해 숫자를 정렬해주어야 한다. 또한, heap은 최소값만 보장해주기 때문에 최소힙과 최대힙을 만들어 사용해야 한다. 최대힙은 힙에 원소를 추가할 때 ‘-‘를 붙여주면 구현할 수 있다. 아래 그림을 보며 이해해보자!

• 우리가 궁금한건 중간값이기 때문에 최소힙과 최대힙을 그림과 같이 동시에 사용해야 한다.
• 최대힙의 최대값을 중간값으로 생각하고 출력할 것이다.
• 원소를 집어넣을 때 원소가 최대힙의 최대값과 같거나 작다면 최대힙에 추가해주고, 아닐경우 최소힙에 추가한다.
• 원소를 추가한 이후 최대힙 원소의 개수가 최소힙 원소의 개수보다 2개 많다면 최대힙에서 최대값을 빼서 최소힙에 집어넣는다.
• 최대힙의 원소의 개수가 최소힙의 원소의 개수보다 적다면 최소힙의 최소값을 최대힙에 집어넣는다.

이렇게 시간복잡도를 줄이며 중간값을 탐지할 수 있다.

Heapq 시간복잡도

• Push : O(log n) **heap 자료형은 push로 원소를 추가할 때 자동으로 정렬된다.
• Pop : O(log n)

List 시간복잡도

• Append : O(1)
• Sort : O(N Log N)

매번 Sort를 통해 정렬하게 된다면 시간복잡도가 heap을 사용했을 때 보다 급격히 증가함을 알 수 있다.

CODE

import sys
import heapq
n = int(input())
up = [] #최소힙
down = [] #최대힙

for _ in range(n):
    num = int(sys.stdin.readline())
    if not down: #최대힙에 원소를 추가하거나 뺄때는 '-'를 붙여주어야 한다.
        heapq.heappush(down, -num)
    else:
        if -down[0] >= num:
            heapq.heappush(down, -num)
        else:
            heapq.heappush(up, num)
    if len(down)-len(up) == 2:
        tmp = -heapq.heappop(down)
        heapq.heappush(up, tmp)
    elif len(down) < len(up):
        tmp = -heapq.heappop(up)
        heapq.heappush(down, tmp)

    print(-down[0])

[다시 풀어보기]