[프로그래머스] Lv3. 이중우선순위큐

🔗

풀이

우선순위 큐를 두 개 구현. 파이썬의 heapq 라이브러리를 사용함.

최대값과 최소값을 둘 다 출력해야 하는데 파이썬의 힙큐 라이브러리는 최소힙이기 때문에 최대힙을 구현하기 위해서는 값을 넣을 때 -를 붙여 수의 부호를 변경해야 한다.

둘 중 어느 하나의 힙에 출력연산이 발생해 값이 출력되어도 다른 힙에는 그 값이 여전히 남아있을 수 있다. 매번 그 값을 찾아서 출력하기는 번거롭기 때문에 나중에 다른 힙에서 그 값을 출력하려고 할 때 그 값이 없다는 사실만 알려줄 수 있도록 digit 딕셔너리를 통해 현재 남아있는 수의 개수를 저장할 수 있도록 했다. 출력 연산이 발생하고 큐가 비어있지 않다면 현재 남아있는 수의 개수가 0이 아닌 수만 출력될 수 있도록 한다.

마지막에 모든 연산이 종료된 후 큐에 값이 남아있지 않을 때는 [0, 0]이 출력될 수 있도록 한다는 점, 큐에 값이 1개만 남아있을 때는 최소값, 최대값이 모두 그 값이라는 예외처리만 따로 해주었다. 남아있는 수의 개수가 2개 이상일 때는 최대힙과 최소힙을 모두 돌려서 가능한 숫자를 하나씩 출력했다.

def solution(operations):
    answer = [0, 0]
    digit = defaultdict(int)
    Q1, Q2 = [], [] # Q1은 최소큐, Q2는 최대큐
                
    for s in operations:
        op, n = s.split()
        n = int(n)
        if op == 'I':
            heappush(Q1, n)
            heappush(Q2, -n)
            digit[n] += 1
        elif n == 1:
            while Q2:
                if digit[-Q2[0]]:
                    digit[-heappop(Q2)] -= 1
                    break
                heappop(Q2)
        elif n == -1:
            while Q1:
                if digit[Q1[0]]:
                    digit[heappop(Q1)] -= 1
                    break
                heappop(Q1)
    
    sum_digit = sum(digit.values())
    if sum_digit == 0:
        return [0, 0]
    
    if sum_digit == 1:
        for k, v in digit.items():
            if v:
                return [k, k]
    
    # 최대 값 반환
    while Q2:
        if digit[-Q2[0]]:
            answer[0] = -heappop(Q2)
            break
        heappop(Q2)
    
    while Q1:
        if digit[Q1[0]]:
            answer[1] = heappop(Q1)
            break
        heappop(Q1)
    return answer

from collections import defaultdict
from heapq import heappop, heappush