BOJ #1655 가운데를 말해요

읽기 전

• 불필요한 코드나 잘못 작성된 내용에 대한 지적은 언제나 환영합니다.
• 개인적으로 배운 점을 정리한 글입니다.

문제 링크

문제 1655. 가운데를 말해요

문제 풀이

시간제한이 0.1초인 걸 봐선 입력과 동시에 중간값을 출력해야 함을 의미한다.

문제 자체는 우선순위 큐로 해결할 수 있는데 중간값이므로 가운데 수보다 작은 값들, 큰 값들을 저장할 우선순위 큐 1개가 아닌 2개로 풀어내야 한다.

전체 숫자가 짝수 개면 가운데 두 수 중 작은 수를 출력하고 그 외엔 중간값을 출력해야 하는데 중간을 기점으로 좌측은 최대 힙을, 우측은 최소 힙에 저장하면 되겠다. 그러나 입력이 미리 주어지지 않기 때문에 매번 중간값을 찾기 위해 적절히 삽입해야 한다. 우선 중간을 기점으로 최대 힙, 최소 힙으로 나누었으니 우선순위 큐의 루트들이 중간 지점을 향하도록 만들었다. 그렇다면 우측의 최소 힙은 좌측 최대힙의 모든 원소보다 값이 커야 한다. 그리고 좌우측 큐는 서로 균형이 맞아야 양쪽 루트가 중간 값임을 보증할 수 있다. 따라서 다음의 규칙을 준수해야 한다.

• 최대 힙, 최소 힙의 길이 차이는 1 이하여야 한다.
• 길이가 부족한 배열에 우선 삽입한다. 같은 경우 최대 힙에 삽입한다.
• 만약 좌측 최대 힙의 루트 노드가 우측 최소 힙의 루트 노드보다 크면 서로 바꾼다.
• 좌측 최대 힙의 루트 노드가 입력 받아온 값들의 중간 값이다. (중간 값이 2개라면 작은 값을 출력하기로 했으므로)

python 코드

import sys
import heapq
input = sys.stdin.readline
 
N = int(input())
left_q = []
right_q = []
for _ in range(N):
    x = int(input())
    if not left_q:
        heapq.heappush(left_q, (-x, x))
        continue
    if len(left_q) <= len(right_q):
        heapq.heappush(left_q, (-x, x))
    elif len(left_q) > len(right_q):
        heapq.heappush(left_q, (x, x))
    if left_q[0][1] > right_q[0][1]:
        l_out = heapq.heappop(left_q)
        r_out = heapq.heappop(right_q)
        heapq.heappush(left_q, (-r_out[1], r_out[1]))
        heapq.heappush(right_q, (l_out[1], l_out[1]))
    print(left_q[0][1])