- 문제 : https://school.programmers.co.kr/learn/courses/30/lessons/42628
문제 이해
- 주어진 연산(operations)을 진행하면서 데이터를 삽입하거나, 최소값 또는 최대값을 제거하는 문제이다.
- 최소값과 최대값은 각각의 heap을 이용하면 된다는 아이디어는 있었지만 하나의 배열에서 진행하기 때문에 최소값을 최소 Heap에서 pop할 때에도 최대 Heap에서 해당 값을 제거해야했다. (최대값을 최대 heap에서 pop할 때에도 마찬가지이다)
- 처음에는 힙에서 원소를 제거하는 방법을 몰라서 while문을 돌면서 최소값 또는 최대값과 동일한 값이 있으면 하나만 빼는 로직을 구현했으나, heap에서는 항상 최소값을 pop하기 때문에 항상 같은 값을 pop한다는 것을 간과했다.
- 그래서 찾아보니 heap에서도 remove를 사용할 수 있다는 것을 알았고, 그것을 적용했다.
- 최대힙, 최소힙을 따로 구현하지 않아도 하나의 힙으로도 진행할 수 있는 방법도 알게 되었는데 방법은 다음과 같다.
heap = heapq.nlargest(len(heap), heap)[1:]
heapq.heapify(heap)nlargest(n, list)를 이용하면 힙에서 가장 큰 수 n개를 뽑아낼 수 있다. 위 코드에서는 heap의 길이만큼의 큰 수를 뽑기에 숫자가 내림차순으로 정렬된다. 이를 인덱스 1부터 슬라이싱 하면 가장 큰 값이 제외될 수 있고, 이를 다시 heapify 하면 된다.
풀이 로직
- 최소힙(heap)과 최대힙(max_heap)을 각각 선언한뒤 while문으로 operations 원소를 pop한다. 이때 원소는 op, number이다.
- operation은 연산이 삽입(I)인지, 삭제(D)인지 알려준다. number는 -1인 경우에는 최소값 삭제, 1인 경우는 최대값 삭제, 그 외는 삽입하는 숫자다.
- 주어진 operation이 I인 경우에는 최소힙에는 주어진 number 값 그대로 push하고, 최대힙에는 (-number, number) 형태로 push한다. (파이썬에서의 heapq는 최소값이 0인덱스인 최대힙이기 때문이다)
- operation이 D이고, numbe 값에 따라 최대 또는 최소를 pop한다.
1) number == -1인 경우는 최소값을 pop해야하므로 최소힙에서는 pop을 하고 해당 값을 최대힙에서 제거(remove)한다.
2) number == 1인 경우는 최대값을 pop해야하기에 최대힙에서 pop 한 뒤 해당 값을 최소 힙에서 제거 한다.
- 반복문이 끝난 뒤 heap에 원소가 없다면 [0, 0]을 원소가 있다면 [최대값, 최소값]을 반환한다.
정답 풀이
from collections import deque
import heapq
def solution(operations):
heap = []
max_heap = []
operations = deque(operations)
while operations:
operation, number = operations.popleft().split()
if operation == 'I':
number = int(number)
heapq.heappush(heap, number)
heapq.heappush(max_heap, (-1 * number, number))
else:
if len(heap) == 0:
continue
if number == '1': # 최댓값 뽑기
x = heapq.heappop(max_heap)[1]
heap.remove(x)
elif number == '-1':
x = heapq.heappop(heap)
max_heap.remove((-1 * x, x))
if len(heap) == 0:
return [0, 0]
else:
return [heapq.heappop(max_heap)[1], heapq.heappop(heap)]
참고