백준 1655 가운데를 말해요 python

처음에는 정렬을 이용해 단순히 구현해 봤는데 시간초과가 나왔다.
이를 힙큐를 사용하여 해결할 수 있었다.

📜 문제 바로 가기 : 가운데를 말해요

제출코드

파이썬

import sys
input = sys.stdin.readline

import heapq

N = int(input())

middleNum = int(input())
smallHQ = []
bigHQ = []
print(middleNum)

for i in range(2, N+1):
    num = int(input())
    if middleNum > num:
        heapq.heappush(smallHQ, -num)
        heapq.heappush(bigHQ, middleNum)
    else:
        heapq.heappush(smallHQ, -middleNum)
        heapq.heappush(bigHQ, num)
    
    if len(smallHQ) >= len(bigHQ):
        middleNum = -1 * heapq.heappop(smallHQ)
    else:
        middleNum = heapq.heappop(bigHQ)

    print(middleNum)

---

접근방법

첫번째 생각 : 그냥 계속 정렬하기

숫자가 들어올 때 마다 배열에 숫자를 넣고 정렬한다.
그리고 지금까지 숫자 개수가 홀수면 arr[i//2], 짝수면 arr[i//2 -1]을 출력한다.

이 로직은 시간초과가 나왔는데 숫자를 받는데 O(N), 정렬에 O(N logN) 의 시간이 걸려 결과적으로는 O(N^2 logN)의 시간이 나와 시간초과가 발생한것 같다.

두번째 생각 : 힙큐로 분할하기

힙큐 2개를 이용한다.
smallHQ, bigHQ 라고 변수명을 지었는데 중간값을 기준으로 smallHQ 에는 보다 작은 값을, bigHQ 에는 보다 큰 값을 넣는다.

smallHQ 에서는 큰값을 꺼내야 되기 때문에 최대힙으로 만들었다.

이제 숫자가 들어올 때 마다 이전까지 중간 값 middleNum과 비교하여 큰값은 bigHQ에, 작은 값은 smallHQ 에 넣는다.

이렇게 하면 둘 중 원소가 더 많은 쪽에서 heappop 을 수행하여 빼면 그 값이 전체에서 중간값이 된다.

전체 개수가 짝수일때는 절반씩 힙큐에 들어가게 되는데 문제에서 그럴때에는 작은값을 뽑으라고 했으므로 smallHQ 에서 뽑으면 된다.

힙큐는 연산에 O(logN)의 시간이 필요하기 때문에 전체 시간은 O(N logN)이 나와 통과핳 수 있었던거 같다.