[PG/Python] 야근 지수

📌 문제

https://programmers.co.kr/learn/courses/30/lessons/12927

코딩테스트 연습 - 야근 지수

 

📌 문제 접근 방법

1. works에 들어있는 일들의 작업량을 key로, 해당 작업량만큼의 작업 개수를 value로 딕셔너리를 만들어줬다.
2. n만큼 반복문을 돌리면서 작업량이 가장 큰 작업부터 1씩 작업을 진행했다.
   이 때, 해당 작업량을 가진 작업이 모두 작업을 진행했다면 해당 key를 삭제하고 가장 큰 작업량을 갱신한다.
   또한 가장 큰 작업량이 0이 된다면 모든 작업을 완료한 것이므로 반복문을 종료한다.
3. n번의 작업이 끝난 후 남은 작업들에 대해 피로도를 계산한다.
   작업량이 0보다 작다면 이미 작업을 완료한 것과 같으므로 계산에 포함하지 않는다.

 

📌 코드

from collections import defaultdict

def solution(n, works):
    works_dict = defaultdict(int)
    for work in works:
        works_dict[work] += 1
        
    max_num = max(works_dict.keys())
    for _ in range(n):
        works_dict[max_num] -= 1
        works_dict[max_num-1] += 1
        
        if works_dict[max_num] == 0:
            del works_dict[max_num]
            max_num = max(works_dict.keys())
        
        if max_num == 0:
            break
            
    answer = 0
    for work, cnt in works_dict.items():
        if work <= 0 :
            continue
        
        for _ in range(cnt):
            answer += work**2
            
    return answer