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

https://school.programmers.co.kr/learn/courses/30/lessons/42628

프로그래머스

 

백준에도 똑같은 문제가 있다.

전에 프로그래머스 레벨 테스트에서 이 문제를 만났다가 구현에서 막혀서 못 풀었던 기억이 났다.

(백준에서도 힘들게 풀었었는데..)

 

이번에 다시 풀어볼 때는 한번에 잘 풀려서 기분이 좋았다. (옛날 풀이는 기억이 나지 않고 완전히 새로운 느낌으로 풀었다.)

이 문제의 핵심은 최소힙과 최대힙을 만들고, 이 둘 사이에 데이터 동기화를 잘 시켜주는 것이다.

 

먼저 데이터를 추가할 때는 두 힙에 모두 데이터를 추가한다.

최솟값을 삭제할 때는 최소힙에서 데이터를 삭제하고, 최댓값을 삭제할 때는 최대힙에서 데이터를 삭제한다.

 

데이터를 삭제할 때마다 deleted 딕셔너리에 삭제한 수를 키로 넣고, 삭제한 횟수를 value 로 저장한다.

그리고, 최솟값을 삭제하기 전에 현재 pop 한 숫자가 다른 힙에서 삭제된 수인지 deleted 딕셔너리를 확인해서 이미 삭제된 수는 빼고 나서 삭제한다.

최댓값도 삭제하기 전에 deleted 배열을 보고 이미 삭제된 수인지 확인해서 동기화 시켜준 뒤 삭제한다.

 

최종적으로 최댓값, 최솟값을 뺄 때도 같은 방식으로 최댓값, 최솟값을 한번씩 뺀다.

이때 이중우선순위큐에 최종 데이터가 1개 남아있을 수도 있는 점만 주의해주면 된다.

 

from heapq import heappush, heappop

def solution(operations):
    max_heap = []
    min_heap = []
    deleted = dict()
    
    for op in operations:
        op, data = op.split()
        data = int(data)
        if op == 'D':
            if data == 1: # delete max
                while max_heap and max_heap[0][1] in deleted and deleted[max_heap[0][1]] > 0:
                    _, popped = heappop(max_heap) # 과거에 삭제된 기록이 있는 데이터 삭제
                    deleted[popped] -= 1
                if max_heap:
                    _, popped = heappop(max_heap)
                    if popped in deleted:
                        deleted[popped] += 1
                    else:
                        deleted[popped] = 1
            else:
                while min_heap and min_heap[0][1] in deleted and deleted[min_heap[0][1]] > 0:
                    _, popped = heappop(min_heap) # 과거에 삭제된 기록이 있는 데이터 삭제
                    deleted[popped] -= 1
                if min_heap:
                    _, popped = heappop(min_heap)
                    if popped in deleted:
                        deleted[popped] += 1
                    else:
                        deleted[popped] = 1
        else:
            heappush(max_heap, (-data, data))
            heappush(min_heap, (data, data))
    
    # pop max
    popped_max = None
    popped_min = None
    while max_heap and max_heap[0][1] in deleted and deleted[max_heap[0][1]] > 0:
        _, popped = heappop(max_heap) # 과거에 삭제된 기록이 있는 데이터 삭제
        deleted[popped] -= 1
    if max_heap:
        _, popped = heappop(max_heap)
        popped_max = popped
        if popped in deleted:
            deleted[popped] += 1
        else:
            deleted[popped] = 1
            
    # pop min
    while min_heap and min_heap[0][1] in deleted and deleted[min_heap[0][1]] > 0:
        _, popped = heappop(min_heap) # 과거에 삭제된 기록이 있는 데이터 삭제
        deleted[popped] -= 1
    if min_heap:
        _, popped = heappop(min_heap)
        popped_min = popped
        if popped in deleted:
            deleted[popped] += 1
        else:
            deleted[popped] = 1
    
    if popped_max == None and popped_min == None:
        return [0, 0]
    
    if popped_max == None:
        popped_max = popped_min
    
    if popped_min == None:
        popped_min = popped_max
    
    return [popped_max, popped_min]

 

중복되는 코드가 많아서 조금 더 깔끔하게 정리하면

 

from heapq import heappush, heappop

def solution(operations):
    max_heap = []
    min_heap = []
    deleted = dict()
    
    def delete_max():
        while max_heap and max_heap[0][1] in deleted and deleted[max_heap[0][1]] > 0:
            _, popped = heappop(max_heap) # 과거에 삭제된 기록이 있는 데이터 삭제
            deleted[popped] -= 1
        if max_heap:
            _, popped = heappop(max_heap)
            if popped in deleted:
                deleted[popped] += 1
            else:
                deleted[popped] = 1
            
            return popped
        return None
    
    def delete_min():
        while min_heap and min_heap[0][1] in deleted and deleted[min_heap[0][1]] > 0:
            _, popped = heappop(min_heap) # 과거에 삭제된 기록이 있는 데이터 삭제
            deleted[popped] -= 1
        if min_heap:
            _, popped = heappop(min_heap)
            if popped in deleted:
                deleted[popped] += 1
            else:
                deleted[popped] = 1
            return popped
        
        return None
    
    for op in operations:
        op, data = op.split()
        data = int(data)
        if op == 'D':
            if data == 1: # delete max
                delete_max()
            else:
                delete_min()
        else:
            heappush(max_heap, (-data, data))
            heappush(min_heap, (data, data))
    
    popped_max = delete_max()
    popped_min = delete_min()
    
    if popped_max == None and popped_min == None:
        return [0, 0]
    
    if popped_max == None:
        popped_max = popped_min
    
    if popped_min == None:
        popped_min = popped_max
    
    return [popped_max, popped_min]

 

그런데 delete 함수에서도 중복되는 코드가 보인다.

더 중복코드를 깔끔하게 분리해보자.

 

from heapq import heappush, heappop

def solution(operations):
    max_heap = []
    min_heap = []
    deleted = dict()
    
    def delete(heap):
        while heap and heap[0][1] in deleted and deleted[heap[0][1]] > 0:
            _, popped = heappop(heap) # 과거에 삭제된 기록이 있는 데이터 삭제
            deleted[popped] -= 1
        if heap:
            _, popped = heappop(heap)
            if popped in deleted:
                deleted[popped] += 1
            else:
                deleted[popped] = 1
            
            return popped
        return None
    
    for op in operations:
        op, data = op.split()
        data = int(data)
        if op == 'D':
            if data == 1: # delete max
                delete(max_heap)
            else:
                delete(min_heap)
        else:
            heappush(max_heap, (-data, data))
            heappush(min_heap, (data, data))
    
    popped_max = delete(max_heap)
    popped_min = delete(min_heap)
    
    if popped_max == None and popped_min == None:
        return [0, 0]
    
    if popped_max == None:
        popped_max = popped_min
    
    if popped_min == None:
        popped_min = popped_max
    
    return [popped_max, popped_min]

 

최종코드는 위와 같다.