[python] heapq 모듈 부시기❗️

프로그래머스 level2의 문제 중 더 맵게 라는 문제를 해결하던 도중에 효율성 테스트를 계속 통과하지 못하여 힙에 대해 생각하던 중,

heapq모듈을 접하면서 쉽게 풀 수 있게 되어 이 모듈에 대해 정리해보고자 한다.

 

 

programmers.co.kr/learn/courses/30/lessons/42626

코딩테스트 연습 - 더 맵게

 

우선 heapq 모듈은 이진트리로 구성된 min heap 구조를 제공해준다.

이를 사용하기 위해서는 아래와 같은 코드를 추가해주어야 한다.

import heapq

 

모듈을 불러온 후, heap 구조를 사용할 수 있게 된다.

 

# 힙을 위한 리스트 생성
heap = []


# 원소 추가 방법
heapq.heappush(heap,2)
heapq.heappush(heap,1)
heapq.heappush(heap,4)
heapq.heappush(heap,3)


# 원소 삭제 방법
heapq.heappop(heap)	#1
# heap = [2,3,4]


# 원소 얻는 방법
print(heap[0])	# 2


# 이미 존재하는 리스트를 힙으로 변환하는 방법
heap = [1,2,3,4,5]
heapq.heapify(heap)

 

기존에 작성하였던 코드는 효율성테스트를 통과하지 못하였다.

이는 while문이 돌아갈 때마다 sort를 하게되어 효율성이 떨어진다고 판단되었다.

def solution(scoville, K):
    answer = 0
    while scoville[0]<=K :
        if len(scoville)==1 :
            return -1
        answer+=1
        scoville.append(scoville.pop(0)+scoville.pop(0)*2)
        scoville.sort()
    return answer

 

heapq를 알게된 후, 다음과 같이 코드를 수정하였는데, 모든 효율성 테스트에서 통과할 수 있게 되었다..!

import heapq
def solution(scoville, K):
    answer = 0
    heapq.heapify(scoville)

    while scoville[0]<=K :
        if len(scoville)==1 :
            return -1
        answer+=1
        heapq.heappush(scoville,heapq.heappop(scoville) + heapq.heappop(scoville)*2)
    return answer

 

시간 복잡도 문제로 어려움을 겪는 경우, heap 사용도 고려해보아야겠다❗️😊