[골드 2] 백준 1655 - 가운데를 말해요 (파이썬)

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

풀이

• max_pq는 왼쪽 절반, 즉 중간값 이하를 담고 최대 힙으로 운용한다. (파이썬에서는 -값을 넣어 사용)
• min_pq는 오른쪽 절반, 즉 중간값 초과를 담고 최소 힙으로 운용한다.
• 입력값과 max_pq, min_pq에서 하나씩 빼서 총 3개의 숫자 후보를 만들고 정렬 후 재분배한다.
• 중간값은 항상 max_pq의 top인 -max_pq[0]이 된다.

max_pq = [1e9]
min_pq = [1e9]

1e9를 미리 넣어둔다. 즉 힙이 비어있을때 heappop 하기 위한 임시값이다.

하지만 이건 비정상적인 초기값이고, 결과적으로 정답이 이상해 질 수 도 있다.

a = [A[i]]
a.append(-heapq.heappop(max_pq))
a.append(heapq.heappop(min_pq))
a.sort()

현재 숫자 + 양쪽 합에서 하나씩 빼서 총 3개로 구성한다.

정렬 후 다시 힙에 재분배한다.

while len(max_pq) < (i + 2) // 2 + 1 :
    heapq.heappush(max_pq, -a[pos])
    pos += 1

while pos < 3 :
    heapq.heappush(min_pq, a[pos])
    pos += 1

max_pq는 (i + 1) // 2개의 원소를 가져야 한다. 즉 중간값이 항상 top에 위치하도록 강제한다.

남은 값은 min_pq로 넘긴다.

문제점은 다음과 같다.

• 1e9를 초깃값으로 힙에 넣어두는 방식은 의도치 않게 잘못된 값이 출력될 수 있다.
• 처음 입력이 1이라도 1e9와 함께 섞이므로 이상한 값 나올수 있다.
• 예를 들어 입력이 [1]이면 출력이 1이어야 하나, 네 코드에서는 1e9까지 섞여서 오작동 할 수 도 있다.

코드

# 백준 1655 - 가운데를 말해요
# 분류 : 우선순위 큐

import sys
import heapq

input = sys.stdin.readline

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

max_pq = [1e9]
min_pq = [1e9]

for i in range(N) :
    a = [A[i]]
    a.append(-heapq.heappop(max_pq))
    a.append(heapq.heappop(min_pq))
    a.sort()

    # i + 1 -> (i + 2) // 2 + 1
    pos = 0
    while len(max_pq) < (i + 2) // 2 + 1 :
        heapq.heappush(max_pq, -a[pos])
        pos += 1
    
    while pos < 3 :
        heapq.heappush(min_pq, a[pos])
        pos += 1

    print(-max_pq[0])