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

백준 1655번 가운데를 말해요 파이썬

문제 링크 - https://www.acmicpc.net/problem/1655

1655번: 가운데를 말해요

N=10만이고 시간 제한이 0.1초인 문제다. 가운데 값이 갱신되는 블랙 박스가 있다고 할 때, 삽입과 조회가 모두 O(log N)안에 해결되어야 한다. B+tree가 생각났는데 이 때부터 계속 트리 구조를 수정하면서 그림을 그렸다. 제대로 구현되는 게 없어서 힌트를 보고 풀었다.

우선순위 큐를 사용하여 가운데에서 가장 가까운 값을 O(1) 시간에 pop할 수 있도록 구현한다. 이를 위해 왼쪽 큐는 최대힙, 오른쪽 큐는 최소힙으로 구현하면 된다.

1. 풀이

• 2개의 우선순위 큐를 사용한다. (오른쪽 큐 - 최소 힙, 왼쪽 큐 - 최대 힙)
• 새롭게 등장한 수가 “가운데"보다 크거나 같다면 오른쪽 큐에, 작다면 왼쪽 큐에 삽입한다.
• 현재까지 등장한 숫자의 개수가 짝수이고 왼쪽 큐의 길이가 더 크다면, “가운데"를 오른쪽 큐에 삽입하고, 왼쪽 큐에서 pop하고 “가운데”로 갱신한다.
• 현재까지 등장한 숫자의 개수가 홀수이고 오른쪽 큐의 길이가 왼쪽 큐의 길이보다 2만큼 크다면, “가운데"를 왼쪽 큐에 삽입하고, 오른쪽 큐에서 pop하고 “가운데"로 갱신한다.

2. 코드

from heapq import heappop, heappush

N = int(input())
leftHeap =[]
rightHeap = []
L = []
for i in range(N):L.append(int(input()))
mid = L[0]
answer=[mid]
for i in range(1,N):
    num = L[i]
    if mid <=num :
        heappush(rightHeap,num)
    else:
        heappush(leftHeap,-1*num)
        
    if (i+1)%2==0 :
        if len(leftHeap) > len(rightHeap) :
            heappush(rightHeap, mid)
            mid = -1*heappop(leftHeap)
            
    else:
        if len(leftHeap)+2==len(rightHeap) :
            heappush(leftHeap, -1*mid)
            mid = heappop(rightHeap)

    answer.append(mid)
       
for a in answer:
    print(a)