[python]programmers - 더맵게

프로그래머스 Level2 힙(Heap) 더 맵게

 

참고 : heapq 모듈 사용법

[파이썬] heapq 모듈 사용법

 

소스 코드

import heapq

def solution(scoville, K):
  h = [] # heap
  c = 0 # 변환 횟수 카운트

  for i in scoville :
    heapq.heappush(h, i)

  while True :
    a = heapq.heappop(h) # 최소값을 pop
    if(a >= K) : # 최소값이 K보다 크면 나머지 값들도 다 K보다 크다
      break
    elif(len(h) == 0 and a < K) : # 최소값 하나 남았는데 K보다 작으면 -1 반환
      return -1
    b = heapq.heappop(h) # 두번째로 작은 값
    new = a + b * 2 # 새로운 스코빌 값
    heapq.heappush(h, new) # 새로운 스코빌값 push
    c += 1 # 변환 횟수 증가
    
  return c

• 처음 풀때는 list로 while문을 돌 때마다 sorted함수를 사용했다. 그랬더니 효율성 테스트에서 다 실패했다
  • heapq 모듈을 사용해서 실행시간을 매우 단축시켰다
• 문제 조건 중 모든 음식의 스코빌 지수를 K이상으로 만들지 못하면 -1을 반환하라는 조건이 있다

풀이

1. heapq로 사용할 리스트 h, 변환횟수를 카운트할 변수 c를 선언했다

2. 입력받은 scoville 지수가 담긴 scoville list를 heapq에 넣어주었다 (for문)

• heaq 원소 삽입 : heapq.heaqpush(리스트, 삽입 값)
• 기존 scoville list를 heap으로 만들어주려면 heapq.heapify(scoville)로 해줘도 된다

3. heapq.heappop을 하게되면 항상 최소값이 튀어나온다 (a = heapq.heappop(h))

• if) 최소값을 pop했는데 K이상이다 => 리스트 안의 값들은 모두 K이상이므로 반복문을 탈출
• elif) 최소값을 pop했는데, 더이상 원소가 남아있지 않고, pop한 원소가 K 미만이다 => 모든 음식의 스코빌지수를 K이상으로 만들 수 없으므로 -1 반환

4. heappop으로 두번째로 작은 값을 pop한다 (b = heapq.heaqpop(h))

5. 새로운 스코빌 지수 값을 만든다 (최소값 + 두번째 최소값 * 2)

6. 새로운 스코빌 지수를 heapq에 넣어준다

7. 변환 횟수를 1 증가시킨다

 

마지막으로 반복문을 나오면 변환횟수를 최종적으로 반환한다

 

heapq

• heapq 모듈은 이진트리 기반의 최소 힙 자료구조를 제공한다
• 최소 힙은 원소들이 항상 정렬된 상태로 추가되고 삭제된다
• 최소 힙에서 가장작은 값은 항상 index 0이다 (이진 트리에서의 root값)