[프로그래머스] 야근 지수

문제 링크

풀이 과정

처음에는 dfs로 풀었는데 테케만 맞고 제출을 눌렀더니 시간 초과가 났다.

answer = 999999999

def dfs(n, works):
    global answer 
    
    if n == 0:
        temp = 0
        for work in works: 
            temp += work ** 2
        answer = min(temp, answer)
        return 
    
    flag = True
    for work in works:
        if work > 0:
            flag = False
            
    if n > 0 and flag:
        temp = 0
        for work in works: 
            temp += work ** 2
        answer = min(temp, answer)
    
    for i in range(len(works)):
        if works[i] == 0:
            continue
        works[i] -= 1
        dfs(n - 1, works)
        works[i] += 1

def solution(n, works):
    global answer
    
    dfs(n, works)
    
    return answer

다음에는 n번 정렬 하면서 가장 큰 값을 1씩 줄이는 방식으로 풀었는데 정답은 다 맞고 효율성에서 틀렸다고 나왔다.

def solution(n, works):
    answer = 0
    
    for i in range(n):
        works.sort(reverse=True)
        works[0] -= 1
    
    for work in works:
        if work < 0:
            continue
        answer += work ** 2
    
    return answer

마지막으로 우선순위큐 알고리즘으로 풀어서 맞았다!✨

import heapq

def solution(n, works):
    if sum(works) <= n:
        return 0
    
    works = [-w for w in works]
    heapq.heapify(works)
    
    for i in range(n):
        ele = heapq.heappop(works)
        ele += 1
        heapq.heappush(works, ele)
    
    return sum([w ** 2 for w in works])

교훈 & 알게 된 것

• 무작정 풀기 보다 시간복잡도 생각하면서 풀기 -> 매번 얻는 교훈이지만 빨리 풀어야 한다는 생각에 냅다 dfs, 완전 탐색을 돌리게 된다..
• python에서는 heapq 라이브러리로 우선순위 큐를 쉽게 사용 할 수 있다. 라이브러리는 최소 힙으로 구현 되어 있지만 요소를 -로 변환 시켜서 최대 힙으로도 사용 할 수 있다.
  예를 들어 [7, 2, 5, 8, 6] 배열이 있을 때 -를 붙여서 heapq 라이브러리를 사용하면 [-8, -7, -6, -5, -2]가 되고, heapq.pop()을 하면 -8이 나오므로 그때 -를 없애면 최대 힙처럼 사용할 수 있다!
• heapq 라이브러리 사용법
  import heapq
  • 주요 함수
    • heapq.heapify(list): list를 힙으로 만들어준다.
    • heapq.heappush(heap, ele): 힙을 유지하면서 ele 추가
    • heapq.heappop(heap): 힙을 유지하면서 가장 작은 요소를 반환
    • heapq.heappushpop(heap, ele): 힙을 유지하면서 ele를 추가하고 가장 작은 요소를 반환
    • heapq.heapreplace(heap, ele): 힙을 유지하면서 가장 작은 요소를 반환하고 ele 추가, 만약 비어있는 힙이라면 IndexError 발생