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

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

프로그래머스

 

처음 시도했을 때 시간초과가 났다.

# 교육 후 모든 신입사원들의 능력치의 합을 최소화
# 오름차순 정렬, 앞 두 자리 더해서 값 바꾸기 (1회)

def solution(ability, number):
    answer = 0
    ability.sort()
    
    for _ in range(number):
        ability.sort()
        power = ability[0] + ability[1]
        ability[0] = ability[1] = power    
    answer = sum(ability)   
    return answer

number가 최대 10,000인데 이렇게 짜면 sort() nlogn 시간복잡도가 10,000 * n * log n이 되어 시간복잡도사 상당히 높아진다.

 

해결방안

heapq 자료구조를 사용한다!

힙 자료구조

-> heapq모듈은 이진 트리 기반의 최소 힙 자료구조를 제공하기 때문에 원소들이 항상 정렬된 상태로 추가되고 삭제된다.

또한, 최소 힙 내에선 모든 원소는 항상 자식 원소들보다 크기가 작거나 같도록 원소가 추가되고 삭제된다.

 

from heapq import heapify, heappush, heappop

def solution(ability, number):
    heapify(ability)
    for _ in range(number):
        a = heappop(ability)
        b = heappop(ability)
        heappush(ability, a+b)
        heappush(ability, a+b)
    return sum(ability)

 

(복습) heapq 모듈 관련 함수 ⭐️

from heapq import heappush, heappop

heap = [] # 최초 힙 생성

heappush(heap, 4) # 힙에 원소 추가
heappop(heap) # 힙에서 원소 삭제 (삭제되면 가장 작은 값이 맨 앞으로 오게 된다. (재정렬))

heap[0] # 최솟값 삭제하지 않고 얻기

li = [1, 9, 0]
heapify(li) # 기존 리스트를 힙으로 전환

nlargest, nsmallest - 힙의 n개의 가장 큰 리스트, n개의 가장 작은 리스트

heap_q = [1, 9, 2, 0, 3]
heapify(heap_q)

heapq.nlargest(n=3, iterable=heap_q) # heap_q에서 가장 큰 3개의 원소가 담긴 리스트
heapq.nsmallest(n=3, iterable=heap_q) # heap_q에서 가장 작은 3개의 원소가 담긴 리스트

 

이렇게 자료구조를 몰라서 못 푸는 문제는 그냥 해설을 빨리 보고 공부하자.