[백준/Python]1655: 가운데를 말해요

문제

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

백준 문제집 - 0x17강 - 우선순위 큐

 

알고리즘: 우선순위 큐

---

풀이

 

현재까지 주어진 수가 홀수개라면 중간값을, 짝수개라면 중간값 중 더 작은값을 출력해야한다.

 

이 문제는 두 개의 힙(left, right)으로 해결할 수 있다.

left에는 중간값보다 작거나 같은, right에는 중간값보다 큰 수들이 들어가게끔 구성한다.

즉 left는 최대 힙으로, right는 최소 힙으로 구현한다.

 

그리고 새로운 수 x가 주어질때마다 left와 right의 길이를 비교한다.

- 두 힙의 길이가 같다면 left에 -x를 삽입 (주어진 수가 짝수개라면 더 작은값을 출력해야하므로)

- 길이가 다르다면 right에 x를 삽입

 

힙에 삽입 후 left의 첫번째 원소값, right의 첫번째 원소값을 비교한다.

- left[0]: 중간값 or 중간값 중 더 작은값

- right[0]: 중간값보다 큰 첫번째 값

(🚨 힙은 완전이진트리이므로 완벽한 정렬은 X, 하지만 0번째 값은 항상 최소값을 보장함)

 

만약 -left[0] 이 right[0]보다 더 크다면 힙을 바꿔 다시 삽입한다.

이렇게 되면 left[0]은 항상 중간값 or 중간값 중 더 작은값 을 보장하게되는것이다!

 

과정을 정리하면 다음과 같다.

1. 새로운 수 x가 입력됨
2. left와 right의 길이를 비교
   1. 길이가 같다면 left에 -x를 삽입 (중간값이 left에 있어야하므로)
   2. 다르다면 right에 x를 삽입
3. -left[0]과 right[0] 비교
   1. -left[0] > right[0] 이라면 pop
   2. left에서 pop한 수 lpop, right에서 pop한 수 rpop
   3. left에 -rpop 삽입, right에 -lpop 삽입
4. -left[0] 출력

 

전체 코드

# 메모리: 37132KB / 시간: 224ms
from sys import stdin
from heapq import heappop, heappush


input = stdin.readline

N = int(input())
left, right = [], []  # left: 최대힙(중간값보다 작거나 같은), right: 최소힙(중간값보다 큰)

for _ in range(N):
    num = int(input())
    
    if len(left) == len(right):
        heappush(left, -num)
    else:
        heappush(right, num)

    if right and -left[0] > right[0]:
        l, r = heappop(left), heappop(right)
        heappush(right, -l)
        heappush(left, -r)

    print(-left[0])