가운데를 말해요

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

입력한 수를 누적하며 가운데 수를 매번 출력해주는 문제입니다.

처음에는 heapq 하나를 생성해 계속 정렬하며 가운데 수를 뱉어내는 방식을 택했습니다.

import heapq

n = int(input())

def heap_middle(nums, k):
  heap = []
  for num in nums:
    heapq.heappush(heap, num)
  middle = nums[0]
  for _ in range(k):
    middle = heapq.heappop(heap)
  return middle

number = []
for _ in range(n):
    number.append(int(input()))
    if (len(number) % 2 == 0):
        length = len(number) // 2 
    if (len(number) % 2 == 1):
        length = len(number) // 2 + 1
    print(heap_middle(number , length))
    

수를 하나 받을 때마다 heapq를 하나 생성해서 넣고 가운데 인덱스까지 뱉어내다 보니 시간복잡도가 비교적 많이 증가해 시간 초과가 나왔습니다. 삽입, 생성 합해서 log(n) + log(n/2)의 시간복잡도가 발생하고, 매번 리스트를 생성해서 돌리다보니 그런 결과가 나온게 아닌가 싶습니다.

그래서 이번에는 heapq를 2개 생성해 middle값 중심으로 나누어 넣는 방식을 택해 문제를 풀었습니다.

import heapq
result = []
N = int(input())
heap1 = []
heap2 = []
for _ in range(N):
    temp_number = int(input())
    if (len(heap1) == len(heap2)):
        heapq.heappush(heap1,(-temp_number, temp_number))
    else:
        heapq.heappush(heap2, temp_number)
    if (len(heap2) > 0 and heap2[0] < heap1[0][1]):
        heap1_element = heapq.heappop(heap1)[1]
        heap2_element = heapq.heappop(heap2)
        heapq.heappush(heap1, (-heap2_element, heap2_element))
        heapq.heappush(heap2, heap1_element)
    result.append(heap1[0][1])
for e in result:
    print(e)
        

다행히 pyp3에서는 시간초과가 발생하지 않아 통과는 했습니다....