가운데를 말해요 - 1655번

우선순위 큐로 풀다가 시간초과로 해결하지 못하여 근본적인 문제를 찾던 도중 힙 2개를 이용하여 풀면 된다는 것을 알게 되었다.

최대힙과 최소힙에 입력받은 숫자들을 절반씩 나눠 저장한 뒤 중간값만 출력하는 방식이다.

 

예를들어 문제에 있는 테스트 케이스를 보면서 사람이 직접 푼다고 가정하자

입력으로 7개의 숫자를 입력받는다. 입력 = [ 1, 5, 2, 10, -99, 7, 5]

입력 : 1 >> 출력 : 1

입력 : 5  >> 출력 : 1

입력 : 2 >> 출력 : 2

입력 : 10 >> 출력 : 2

입력 : -99 >> 출력 : 2

입력 : 7 >> 출력 : 2

입력 : 5 >> 출력 : 5

이렇게 표로 표현해볼 수 있다. 

여기서 최대힙과 최소힙을 사용하여 입력받는 수들을 2부분으로 쪼갤 것이다.

처음부터 같은 테스트 케이스를 입력받는데,

어떤 방식으로 최대힙과 최소힙이 나눠져야 할지 표로 나타내보겠다.

입력 : 1 >> 출력 : 1

입력 : 5  >> 출력 : 1

입력 : 2 >> 출력 : 2

입력 : 10 >> 출력 : 2

입력 : -99 >> 출력 : 2

입력 : 7 >> 출력 : 2

입력 : 5 >> 출력 : 5

규칙은 최대힙 크기가 먼저 값을 1개 넣고 그 다음 최소힙에 넣으면서 번갈아 가면서 넣는다.

여기서 최대힙의 인덱스 0번째의 값은 가장 큰 값이,

최소힙의 인덱스 0번째의 값은 가장 작은 값이 들어가게 되는데

최대힙의 0번째 값과 최소힙의 0번째 값을 비교하여 더 작은 수를 최대힙의 0번째로 만들어 주면서 중간값을 수정해 나가야한다.

최대힙은 입력받은 수들 중 작은 수들이 들어가게되고, 최소힙은 큰 수들이 들어가게되어서

인덱스 0번째의 값들이 정답에 유력한 중간값들의 후보가 되는 것이다.

 

# 정답코드
import heapq
import sys

t = int(sys.stdin.readline())
max_heap = []
min_heap = []
for _ in range(t):
    inn = int(sys.stdin.readline())
    if len(max_heap) == len(min_heap):
        heapq.heappush(max_heap, (-1*inn, inn))
    else:
        heapq.heappush(min_heap,(inn,inn))
    
    if min_heap and min_heap[0][1] < max_heap[0][1]:
        too_big = heapq.heappop(max_heap)[1]
        this_is_answer = heapq.heappop(min_heap)[1]

        heapq.heappush(max_heap, (-this_is_answer, this_is_answer))
        heapq.heappush(min_heap, (too_big, too_big))
    # print(min_heap, max_heap)
    # print('ans :',max_heap[0][1])
    print(max_heap[0][1])

 

첫번째로, 우선순위큐 모듈을 이용하여 구현해봤지만 시간이 너무 오래걸린다.

# 우선순위큐를 이용하여 풀었지만 지저분하게 풀어서 시간초과가 나버림
from queue import PriorityQueue
import sys


t = int(sys.stdin.readline())

mainQue = PriorityQueue(maxsize = t)
que  = PriorityQueue(maxsize = t)
que_arr = []
for i in range(t):
    inn = int(sys.stdin.readline())
    mainQue.put(inn)
    que_arr.append(inn)

    for j in range(len(que_arr)):
        que.put(que_arr[j])
    if que.qsize()%2==1:
        for k in range(que.qsize()//2):
            que.get()
        print(que.get())
    else:
        for k in range(que.qsize()//2-1):
            que.get()
        print(que.get())
                
    while not que.empty():
        que.get()