[프로그래머스] 야근지수(Java/Python)

문제 설명

회사원 Demi는 가끔은 야근을 하는데요, 야근을 하면 야근 피로도가 쌓입니다. 야근 피로도는 야근을 시작한 시점에서 남은 일의 작업량을 제곱하여 더한 값입니다. Demi는 N시간 동안 야근 피로도를 최소화하도록 일할 겁니다.Demi가 1시간 동안 작업량 1만큼을 처리할 수 있다고 할 때, 퇴근까지 남은 N 시간과 각 일에 대한 작업량 works에 대해 야근 피로도를 최소화한 값을 리턴하는 함수 solution을 완성해주세요.

제한사항

• works는 길이 1 이상, 20,000 이하인 배열입니다.
• works의 원소는 50000 이하인 자연수입니다.
• n은 1,000,000 이하인 자연수입니다.

전체코드(Java)

import java.util.PriorityQueue;
import java.util.Collections;

class Solution {
    public long solution(int n, int[] works) {
        // 우선순위 Queue를 생성하여 크기가 역순으로 정렬되도록 만든다.
        PriorityQueue<Integer> overtime = new PriorityQueue<>(Collections.reverseOrder());
        // works의 시간을 하나씩 추가한다.
        for(int work : works){
            overtime.offer(work);
        }
        
        // 맨 앞에 있는 값을 빼서 0이 아니라면 1씩 감하여 다시 넣는다.
        for(int i=0;i<n;i++){
            int max = (int)overtime.poll();
            if(max<=0) break;
            overtime.offer(max-1);
        }
        
        return sumPow(overtime);
    }
    
    // 우선순위 queue에 있는 work를 제곱한 값을 sum에 누적하여 반환한다.
    public long sumPow(PriorityQueue<Integer> works){
        long sum = 0;
        while(!works.isEmpty()){
            sum += Math.pow(works.poll(),2);
        }
        return sum;
    }
}

전체코드(Python)

import heapq
def solution(n, works):
    answer = 0
    # heaq이 가장 작은것부터 정렬하므로 works의 값들을 음수화 시켜서 가장 큰 값이 앞으로 갈 수 있도록 변경한다.
    works = [-i for i in works]
    # list화 된 값들을 heap으로 변경한다.
    heapq.heapify(works)
    # 야근을 하는 시간만큼 순회하며 가장 작은값에 하나씩 더하되 0일경우 더하지 않는다.
    for i in range(n):
        num = heapq.heappop(works)    
        num = num+1 if num < 0 else 0
        heapq.heappush(works,num)
    
    # 모두 계산한 값에 제곱늘 취하여 합을 구한다.
    answer = sum(list(map(lambda x:x**2,works)))
    
    return answer

알아야 할 함수들

• Java
  • PriorityQueue(우선순위 큐)
    • 우선순위를 결정하고 우선순위가 높은 엘리먼트가 먼저 나가는 자료구조입니다.
    • 우선순위 힙을 이용하여 구현하는것이 일반적입니다.
    • 데이터 삽입할 때 우선순위를 기준으로 최대힙 혹은 최소  힙을 구성하고 데이터를 꺼낼 때 루트 노드를 얻어낸 뒤 루트 노드를 삭제할 때 는 빈 루트 노드 위치에 맨 마지막 노드를 삽입한 후 아래로 내려가면서 적절한 자리를 찾아서 옮기는 방식으로 진행됩니다.
    • PriorityQueue 선언 : PriorityQueue<Integer> overtime = new PriorityQueue<>(Collections.reverseOrder());[ 우선순위 큐의 자료형과 정렬 기준을 정한다);
    • PriorityQueue에 원소 추가 : overtime.add(work); / overtime.offer(work);
      add를 사용할 경우 삽입에 성공했으면 true, 실패했으면 IllegalStateException을 발생시킨다.
    • PriorityQueue에 원소 삭제 : overtime.poll(); / overtime.remove();
      poll을 사용할경우 첫번째 값을 반환하고 제거(비어 있다면 null을 반환)
    • PriorityQueue에 가장 우선순위가 높은 값 출력 : overtiem.peek();
      PriorityQueue에 가장 우선순위가 높은 값을 반환한다.
•  Python
  •  heapq
    • 이진트리(binary tree) 기반의 최소 힙(min heap) 자료구조를 제공합니다. 자바에 익숙한 분이라면 priorityQueue 클래스라고 생각하시면 쉬울것 같습니다.
    • min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, min heap에서 가장 작은값은 언제나 인덱스 0, 즉 이진 트리의 루트에 위치합니다.
    • 힙에 원소 추가 : heapq.heappush(리스트,추가할 원소)[ 추가시 이진트리의 자료구조에 맞게 추가된다 ]
    • 힙에서 원소 삭제 : heapq.heappop(리스트)[ 항상 작은 원소(맨 앞의 요소)가 제거된다 ]
    • 기존 리스트를 힙으로 변환 : heap.heapify(리스트)

[참고]

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

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

https://coding-factory.tistory.com/603

[Java] PriorityQueue(우선순위 큐) 클래스 사용법 & 예제 총정리

https://www.daleseo.com/python-heapq/

[파이썬] heapq 모듈 사용법