(우선순위큐) 백준 1655번 _ 가운데를 말해요 [Python]

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

1655번: 가운데를 말해요

1655번 문제 : 가운데를 말해요

 

이 문제는 우선순위큐 알고리즘으로 해결해야 하는 문제이다.

처음 문제를 접했을때 큐를 사용해야겠다는 생각은 했으나, 중앙값을 리턴하는 메소드를 구현할때 sorted()를 이용하였다.

파이썬에서 sorted()를 사용하면 대다수의 경우에서 시간초과가 계속 나타나는데 역시나 시간초과로 해결하지 못하였다.

아래 코드는 내가 처음 시도했던 코드이다.

 

import sys
class myQueue(object):
    def __init__(self):
        self.queue = []
    def enqueue(self, n):
        self.queue.append(n)
        pass
    def sort(self):
        self.queue = sorted(self.queue)
        
    def median(self):
        length = len(self.queue)
        if length % 2 == 0:
            return self.queue[int(length/2)-1]
        return self.queue[int(length-1)//2]
    
N = int(sys.stdin.readline())
q = myQueue()

for _ in range(N):
    q.enqueue(int(sys.stdin.readline()))
    q.sort()
    print(q.median())

참고로 input() 대신 sys.stdin.readline()을 사용하면 시간초과를 어느정도 예방할 수 있다고 한다.

 

아무튼, 그렇다면 이 문제에서 중앙값은 sorted()를 사용하지 않고 어떻게 해결할 수 있을까?

이때 필요한 것이 바로 우선순위 큐(Priority Queue)다.

우선순위 큐는 모든 원소에 우선순위가 있고, 이 우선순위가 높은 원소부터 제거하는 자료구조를 의미한다.

 

1655번 설명

위 알고리즘을 코드로 구현하면 아래와 같다.

# import sys/
import heapq


if __name__ == "__main__":
    n = int(input())
    max_h, min_h = [], []

    # max_h[0][1]값을 기준으로 큰 값은 min_h, 같거나 작은 값은 max_h에 삽입
    for _ in range(n):
        num = int(input())
        if len(max_h) == len(min_h):
            heapq.heappush(max_h, (-num, num))
        else:
            heapq.heappush(min_h, (num, num))
        if len(max_h) >= 1 and len(min_h) >= 1 and max_h[0][1] > min_h[0][1]:
            max_value = heapq.heappop(max_h)[1]
            min_value = heapq.heappop(min_h)[1]
            heapq.heappush(max_h, (-min_value, min_value))
            heapq.heappush(min_h, (max_value, max_value))
        print('num is',num)
        print(max_h)
        print(min_h)
        print(max_h[0][1])

참고로 파이썬의 heapq는 기본적으로 최소힙이기 때문에 최대힙(left)에 삽입해 줄때는, -를 붙여서 -num으로 넣어준다.