[python] 프로그래머스 - 이중우선순위큐

출처 : 프로그래머스, https://school.programmers.co.kr/learn/courses/30/lessons/42628

프로그래머스

 

백준에도 동일한 문제가 있다.

 

# I 숫자 -> 큐에 숫자 삽입
# D 1    -> 최댓값 삭제
# D -1   -> 최솟값 삭제
# 큐가 비어있다면 삭제 명령 무시

# [최댓값, 최솟값] return
# 만약 큐가 비어있다면 [0, 0] return

 

 

먼저, 생각해야 할 부분

 

1. min_heap, max_heap 2개가 필요하다는 것까진 생각했을 것이다.

     그렇다면 max_heap에서 뽑아서 버린 최댓값을 min_heap에서 어떻게 제거할 수 있을까?

    

    정확히 말하면, min_heap에서 제거할 필요는 없다.

    핵심은, min_heap에게 max_heap에서 뽑은 데이터가 필요 없다고 알리는 것이다.

    이를 위해서는 하나의 값을 두 heap이 공유하도록 만들어줄 필요가 있다.

---

 

풀이

 

from heapq import heappush, heappop

INSERT = 'I'
DELETE = 'D'

DELETE_MAX = 1
DELETE_MIN = -1

def solution(operations):
    answer = []

    min_heap = []
    max_heap = []

    # operation 실행
    for operation in operations:
        oper, num = operation.split()
        num = int(num)

        if oper == INSERT:
            # 삽입 시, Node를 생성해서, min, max 힙에 넣는다.
            new_node = [num, True]
            heappush(min_heap, (num, new_node))
            heappush(max_heap, (-num, new_node))

        else:  # oper == DELETE
            if num == DELETE_MAX:
                # valid한 노드를 꺼낼 때까지 반복한다.
                # heap의 길이가 0보다 클 동안

                while len(max_heap) > 0:
                    order, pop_node = heappop(max_heap)
                    pop_value, pop_valid = pop_node

                    if pop_valid:
                        # valid하다면, invalid로 바꾼다.
                        pop_node[1] = False
                        break
            elif num == DELETE_MIN:
                while len(min_heap) > 0:
                    order, pop_node = heappop(min_heap)
                    pop_value, pop_valid = pop_node

                    if pop_valid:
                        pop_node[1] = False

                        break

    # answer 갱신

    # 최댓값 추가
    # while로 빼면서 valid한 값이 나오면 추가
    # 최댓값이 있다면 무조건 최솟값도 있음

    while len(max_heap) > 0:
        order, pop_node = heappop(max_heap)
        pop_value, pop_valid = pop_node

        if pop_valid and len(answer) == 0:
            answer.append(pop_value)
            break

    if len(answer) == 0:
        return [0, 0]

    # 최솟값 추가
    # while로 빼면서 valid한 값이 나오면 추가

    while len(min_heap) > 0:
        order, pop_node = heappop(min_heap)
        pop_value, pop_valid = pop_node

        if pop_valid:
            answer.append(pop_value)
            break

    return answer

---

 

알게 된 점

 

heapq 모듈의 커스텀 정렬

 

처음에는 아래와 같은 노드 클래스를 선언해서 문제를 풀려고 했다.

 

class Node:
    def __init__(self, value):
        self.value = value
        self.valid = True
        
def solution(operations):

    ....
    
    new_node = Node(num)

    heappush(min_heap, (num, new_node))
    heappush(max_heap, (-num, new_node))
    
    ....

 

그런데 계속 프로그래머스에서는 테스트케이스 5번이 틀리고 백준에서는 TypeError가 나오길래, 

TypeError가 나올 부분이 대체 어디인지 고민하다가

Node라는 객체를 heapq 모듈이 대소비교를 못해서 생긴 문제라는 것을 깨달았다.

 

힙 정렬의 우선순위가 num을 비교하고 이것이 같을 때, new_node를 < 연산자를 통해 비교하려 드는데

Node 객체를 < 연산자를 통해 비교할 수가 없어 TypeError가 나는 것이다.

 

그래서...굳이 클래스를 쓰고 싶다면 < 연산자를 위한 메소드인 __lt__ 메소드를 오버라이딩 해줄 필요가 있다.

 

class Node:
    def __init__(self, value):
        self.value = value
        self.valid = True

    def __lt__(self, other):
        if self.value >= other.value:
            return True
        else:
            return False

 

다만 라이브러리 내부적으로 어떤 연산자를 사용해서 동작하는지 다 아는 것은 불가능하기 때문에

실제 코딩테스트에서는 클래스를 쓰는 것을 지양하고,

리스트나 튜플, 딕셔너리와 같은 가변 객체를 사용하는 방향으로 작성해야 맞다.

 

위와 같은 클래스를 가장 잘 대체할 수 있는 기본 자료형은 딕셔너리 자료형이다.

이를 이용하면 다음과 같이 수정할 수 있다.

 

VALUE = 'value'
VALID = 'valid'

def solution(operations):
    
    ....
    
    new_node = dict()
    new_node[VALUE] = num
    new_node[VALID] = True

    heappush(min_heap, (num, new_node))
    heappush(max_heap, (-num, new_node))
    
    ....

 

딕셔너리 자료형은 기본 자료형이므로, heapq와 같은 모듈에 대해서도 정상적으로 대소비교가 동작한다.

 

Github

GitHub - bh2980/Algorithm-Problem: 알고리즘 풀이 흔적들

 

Reference

파이썬 heapq 커스텀 정렬 이용하기