[프로그래머스/파이썬] (힙(Heap)) 이중우선순위큐

출처

문제 설명

이중 우선순위 큐는 다음 연산을 할 수 있는 자료구조를 말합니다.

명령어	수신 탑(높이)
I 숫자	큐에 주어진 숫자를 삽입합니다.
D 1	큐에서 최댓값을 삭제합니다.
D -1	큐에서 최솟값을 삭제합니다.

이중 우선순위 큐가 할 연산 operations가 매개변수로 주어질 때, 모든 연산을 처리한 후 큐가 비어있으면 [0,0] 비어있지 않으면 [최댓값, 최솟값]을 return 하도록 solution 함수를 구현해주세요.

제한사항

• operations는 길이가 1 이상 1,000,000 이하인 문자열 배열입니다.

• operations의 원소는 큐가 수행할 연산을 나타냅니다.

  • 원소는 “명령어 데이터” 형식으로 주어집니다.
  • 최댓값/최솟값을 삭제하는 연산에서 최댓값/최솟값이 둘 이상인 경우, 하나만 삭제합니다.

• 빈 큐에 데이터를 삭제하라는 연산이 주어질 경우, 해당 연산은 무시합니다.

입출력 예

operations	return
["I 16","D 1"]	[0,0]
["I 7","I 5","I -5","D -1"]	[7,5]

입출력 예 설명

16을 삽입 후 최댓값을 삭제합니다. 비어있으므로 [0,0]을 반환합니다.
7,5,-5를 삽입 후 최솟값을 삭제합니다. 최대값 7, 최소값 5를 반환합니다.

내 풀이

from heapq import heappop, heappush

def solution(operations):
    arr = []
    for i in range(len(operations)):
        current_op = operations[i]
        if current_op[0] == "I":
            heappush(arr, int(current_op[2:]))
        elif current_op[2:] == "1" and arr:
            arr.remove(max(arr))
        elif arr:
            heappop(arr)
    if arr == []:
        return [0,0]
    else:
        return [max(arr), arr[0]]

for문에서 operations에서 연산을 순서대로 꺼내서 current_op에 대입한다.

연산이 삽입인지 아니면 삭제인지는 current_op의 첫 글자를 보면 알 수 있다.
만약에 current_op[0]이 "I"라면 삽입 연산이기 때문에 세 번째 글자부터 끝까지에 해당하는 숫자인 current_op[2:]를 arr에 넣어야 한다.

current_op[0]이 "I"가 아니면 삭제연산이다.
current_op[2:]가 "1"이면 큐에서 최댓값을 삭제해야 한다.
하지만 큐가 비어있다면 연산은 무시되므로 elif문에 이 두 조건을 모두 넣어주고 elif문 안에서 arr의 최댓값을 remove한다.
이 작업은 heap 자료구조에 영향을 미치지 않는다.

그 외에는 최솟값을 삭제하는 연산일 수 있는데 이 때도 마찬가지로 빈 큐면 연산이 무시되므로 큐가 비었는지만 확인한 뒤 비어있지 않다면 arr에서 heappop을 수행한다.
heap 자료구조에서 heappop을 하면 최솟값을 pop하되 이후의 리스트도 heap 자료구조를 유지한다.

for문에서 빠져나오면 if문을 통해 비어있다면 [0,0]을 반환하고 비어있지 않다면 [최댓값,최솟값]을 반환한다.
마찬가지로 heap 구조의 최솟값은 arr[0]임을 이용했다.

정확성 테스트: ~0.03ms / 10.4MB

---

사실 이 문제는 heapq 모듈을 사용하지 않고도 다음과 같이 풀 수 있다.

def solution(operations):
    arr = []
    for i in range(len(operations)):
        current_op = operations[i]
        if current_op[0] == "I":
            arr.append(int(current_op[2:]))
        elif current_op[2:] == "1" and arr:
            arr.remove(max(arr))
        elif arr:
            arr.remove(min(arr))
    if arr == []:
        return [0,0]
    else:
        return [max(arr), arr[0]]