문제
https://school.programmers.co.kr/learn/courses/15009/lessons/121688
프로그래머스
코드 중심의 개발자 채용. 스택 기반의 포지션 매칭. 프로그래머스의 개발자 맞춤형 프로필을 등록하고, 나와 기술 궁합이 잘 맞는 기업들을 매칭 받으세요.
programmers.co.kr
풀이
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)

'Coding Test > Python' 카테고리의 다른 글
| [프로그래머스] 디스크 컨트롤러 (1) | 2024.03.31 |
|---|---|
| [백준] 2133 타일 채우기 (0) | 2024.03.27 |
| [PCCP 모의고사 2] 3번 (0) | 2024.03.13 |
| [PCCP 모의고사 #2] 1번 - 실습용 로봇 (1) | 2024.03.12 |
| [PCCP 모의고사 #1] 3번 - 유전법칙 (0) | 2024.03.04 |