[BOJ] 1655번 가운데를 말해요

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

풀이소요시간: 1시간

수가 계속해서 들어올 때 가운데 수를 구하는 문제이다. 그런데 시간 제한이 0.1초여서 제한이 있다. 매 입력마다 정렬된 리스트를 구해서 가운데 값을 찾으면 $O(N^2)$으로 시간초과가 발생한다(N<=100,000)

이때, 우리는 가운데값만 찾으면 된다는데에 집중해보면 시간복잡도를 낮출 수 있다. 나는 가운데 값 양쪽에 Priority Queue를 둬서 $O(NlogN)$으로 접근했다.

총 수가 홀수일때, 양쪽 PQ의 원소수는 같아야 하고, 총 수가 짝수일때는 오른쪽 PQ의 원소수가 왼쪽 PQ보다 하나 더 많아야 한다(문제 조건에서 가운데 값이 2개일때 작은 값이 가운데 값이라고 정의해서)

새로운 값이 들어왔을 때
1. 우선 가운데 값 기준으로 left_pq 또는 right_pq 중 한 군데에 새로운 값 넣기
2. left_pq, right_pq의 원소수가 균일하게 분배되게 하기
3. 가운데 값 2개일 때 작은 값을 가운데 값으로 가지게 하기

import heapq

N = int(input())
nums = [int(input()) for _ in range(N)]

middle_num = nums[0]
left_pq = [] # max heap
right_pq = [] # min heap

print(middle_num)

for i, num in enumerate(nums[1:]):
  num_count = i+2 # 현재 모든 수(i=0, 수는 2개)
  
  if middle_num < num:
    heapq.heappush(right_pq, (num,num))
  else:
    heapq.heappush(left_pq, (-num,num))

  if len(left_pq) >= len(right_pq) + 2: # 왼쪽이 더 많을때
    heapq.heappush(right_pq, (middle_num, middle_num))
    middle_num = heapq.heappop(left_pq)[1]
  elif len(left_pq) + 2 <= len(right_pq): # 오른쪽이 더 많을 때
    heapq.heappush(left_pq, (-middle_num, middle_num))
    middle_num = heapq.heappop(right_pq)[1]

  # 총 수가 짝수개일때, 왼쪽이 더 많으면(총 수가 홀수개면 양쪽 원소수 위에서 같아짐)
  if len(left_pq) > len(right_pq):
    heapq.heappush(right_pq, (middle_num, middle_num))
    middle_num = heapq.heappop(left_pq)[1]
    
  print(middle_num)