본문 바로가기

알고리즘

99클럽 코테 스터디 10일차 TIL

문제 설명



풀이 과정 & 코드

  • heap을 사용해서 문제를 풀어야 하는걸 직감함. 하지만 python에서는 heapq가 최소만 구현됨
  • 따라서 최대값도 구현하기 위해서는 최소힙을 가지고 '-'를 붙혀 최대힙을 만들어야 함
  • 수가 들어올 때는 최소힙은 그냥 push하고 최대힙은 수를 음수로 뒤집어서 집어넣으면 항상 최대값이 제일 root에 있게 됨
  • 삭제할 때는 최대, 최소힙에서 수를 둘다 제거해야하는데 최소값을 제거하면 pop한 값을 최대힙에서도 remove를 통해 제거하고 heappop으로 제거한게 아니기 때문에 힙을 유지하지 못해서 heapify로 다시 힙 구조를 유지해야함. 최대값을 제거할 때도 마찬가지
  • 큐가 2개이기 떄문에 count변수를 사용해서 넣고 뺄때 증감을 통해 0이 되면 최소, 최대힙에 수가 비었음을 뜻하기 때문에 두 리스트 둘 다 empty로 초기화를 해야함
  • 이후 for loop이 끝나고 count변수가 0이라면 [0, 0]을 return하고 아니라면 각각의 힙에서 최대, 최소값을 리스트에 담아서 return 하면 됨
def solution(operations):
    import heapq
    answer = [0, 0]
    max_heap, min_heap = [], []
    cnt = 0
    for op in operations:
        order, num = op.split()
        if order == 'I':
            heapq.heappush(max_heap, -int(num))
            heapq.heappush(min_heap, int(num))
            cnt += 1

        elif order == 'D':
            if min_heap and int(num) == -1:
                min_ = heapq.heappop(min_heap)
                max_heap.remove(-min_)
                heapq.heapify(max_heap)
            elif max_heap and int(num) == 1:
                max_ = heapq.heappop(max_heap)
                min_heap.remove(-max_)
                heapq.heapify(min_heap)
            cnt -= 1

        if cnt == 0:
            max_heap, min_heap = [], []

    if cnt == 0:
        return answer

    answer[0], answer[1] = -heapq.heappop((max_heap)), heapq.heappop(min_heap)
    return answer