[PCCP 모의고사 #2] 2번 - 신입사원 교육

문제

https://school.programmers.co.kr/learn/courses/15009/lessons/121688

프로그래머스

 

---

풀이

import heapq

def solution(ability, number):
    q = []
    for a in ability:
        heapq.heappush(q, a)

    for _ in range(number):
        x, y = heapq.heappop(q), heapq.heappop(q)
        new = x + y
        heapq.heappush(q, new)
        heapq.heappush(q, new)

    return sum(q)

 

---

설명

- 매번 가장 작은 2명의 신입 사원을 뽑아야 함. -> 정렬 해서 최솟값 2개를 사용해야 한다!

근데 이때 일반 정렬 sort() 를 사용해서 했더니 시간 초과가 발생했다. 

이를 해결하기 위해 heapq,  즉 최댓값과 최솟값을 빠르게 찾아주기 위해 고안된 완전이진트리를 사용한다.

 

** heapq (priority queue)

- 모든 부모 노드는 자식 노드보다 값이 작거나 큰 이진트리 구조인데, 내부적으로는 인덱스 0에서 시작해 k번째 원소가 항상 자식 원소들(2k+1, 2k+2) 보다 작거나 같은 최소 힙의 형태로 정렬된다.

 

- heapq.heappush(heap, item): item 을 heap에 추가

- heapq.heappop(heap): heap에서 가장 작은 원소를 pop 한 뒤 리턴. 비어있는 경우 IndexError 호출됨.

- heapq.heapify(x): 리스트 x를 즉각적으로 heap으로 변환함. (O(N))

import heapq

heap2 = [50, 10, 20]
heapq.heapify(heap2)

print(heap2)

 

** 최대 힙 만들기

파이썬의 heapq 모듈은 최소 heap 으로 구현되어 있기 때문에 최대 heap 구현을 위해서는 트릭이 필요하다.

아이디어: y=-x 변환을 하면 최솟값 정렬이 최댓값 정렬로 바뀐다!

 

힙에 원소를 추가할 때 (-item, item)의 튜플 형태로 넣어주면 튜플의 첫 번째 원소를 우선순위로 힙을 구성하게 된다.  이때 원소 값의 부호를 바꿨기 때문에, 최소 힙으로 구현된 heapq 모듈을 최대 힙 구현에 활용하게 되는 것이다.

아래의 예시는 리스트 heap_items에 있는 원소들을 max_heap이라는 최대 힙 자료구조로 만드는 코드이다.

heap_items = [1,3,5,7,9]

max_heap = []
for item in heap_items:
  heapq.heappush(max_heap, (-item, item))

print(max_heap)