[백준]_1655번 : 가운데를 말해요(파이썬)

1655번: 가운데를 말해요

 

문제를 보면 정렬되지 않은 데이터가 입력되고 입력이 일어날 때마다 정렬을 한 후에 가운데 요소를 출력해야 한다.

 

하지만 입력이 있을 때마다 데이터들을 정렬하는 것은 시간 복잡도가 $O(N^2log{N})$이므로 바람직하지 않다. 특히나 이문제는 시간제한이 0.1초이므로 매우 빠듯한 편에 속한다. 

 

데이터 입력이 일어날 때마다 정렬을 해주는 우선순위 큐(heapq)를 이용해주자. 힙큐는 원소의 추가(heappush)가 $O(logN)$, 원소 추출(heappop)이 $O(logN)$인 이 문제를 품에 있어서 아주 유용한 모듈이다.

 

하지만 그렇다고 힙큐를 적용한 리스트가 우리가 원하는 방식으로 정렬되는 것이 아니라 heappop을 할 때만 원소들이 오름차순으로 나올 뿐이다. 즉 heappush를 한 리스트(ls)에 대해 ls[len(ls)//2]를 한다고 해서 우리가 원하는 답이 나오지 않는다. 힙큐에 의해 heapify 된 리스트에서는 오직 맨 앞(heapq.heappop(ls) == ls[0] == min(ls)는 항상 참이다!!)만이 유의미하다. 어차피 우리가 보고 싶은 것은 정렬된 리스트의 중간 부분이다. 즉 힙큐를 적용한 리스트를 2개 만들어서 데이터를 관리하자.

 

전략

 

결국 우리가 원하는 이상적인 상황은 이것이다. left, right 힙큐가 2개있다. 두 힙큐에 왼쪽부터 시작해서 번갈아가며 데이터가 들어갈 것이다. 즉 이제까지 들어온 데이터에 대해 중앙값에서 작은 부분들은 left, 큰 부분들은 right로 들어가게 하는 것이다. 데이터 개수가 홀수든 짝수든 left에서 pop 하면 된다.

 

이때 left는 최대힙이고 right는 최소힙이다. 그래야지만 heappop을 적용하면 전체 데이터의 중간 부분(우리가 답으로 관찰하고 싶은 부분)들이 튀어나올 것이기 때문이다.  

 

힙큐 left, right 시각화

 

여기서 오해하면 안되는 것이 left는 최대힙이기 때문에 기존 값의 부호를 바꾼 상태로 입력된다. 하지만 시각화 그림에서는 그냥 부호 반전 부분을 제외하고 개념적인 측면에서만 그렸다.

 

이제 기반 코드들을 작성해보자

import heapq
import sys
input = sys.stdin.readline

left, right = [], []
for i in range(int(input())):
    num = int(input())
    if i%2: # left
        heapq.heappush(right, num)
    else: # right
        heapq.heappush(left, -num)

 

왼쪽 오른쪽 번갈아가면서 값을 넣어주는 것까지는 별 문제없어 보인다. 하지만 아래의 예시를 봐보자.

 

마지막 정답은 2여야한다.

 

이러한 예외를 처리해 주기 위해 아래의 과정이 추가로 필요하다.

 

맨위의 원소들만 비교해서 자리를 바꿔줌

전체 코드 >>

import heapq
import sys
input = sys.stdin.readline

left, right = [], []
for i in range(int(input())):
    num = int(input())
    if i%2: # left
        heapq.heappush(right, num)
    else: # right
        heapq.heappush(left, -num)
    if right and -left[0] > right[0]:
        r = heapq.heappop(right)
        l = -heapq.heappop(left)
        heapq.heappush(right, l)
        heapq.heappush(left, -r)
    print(-left[0])