[Lv.3] 야근 지수

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

class Solution {
    public long solution(int n, int[] works) {
        long answer = 0;
        PriorityQueue<Integer> pq=new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2-o1;
            }
        });
        for (int work : works) {
            pq.add(work);
        }
        while(true){
            if(n==0){
                for (Integer integer : pq) {
                    answer+=integer*integer;
                }
                break;
            }
            Integer target = pq.poll();
            if(target<1){
                return 0;
            }
            pq.add(--target);
            n--;
        }
        return answer;
    }
}

 

 

PriorityQueue를 이용해 쉽게 풀 수 있는 문제였다.

남은 작업량의 제곱의 최소를 만들기 위해선 작업 배열의 최대값을 최소로 만들어야 한다.

즉, 큰 작업부터 먼저 해야한다.

 

로직은 단순하다.

1. 남은 작업시간이 없다면 (퇴근 시간이 되었다) 남아있는 작업의 제곱의 합을 리턴한다.

2. 남은 작업시간은 있으나 할 작업이 없는 경우 (pq의 Poll값이 1보다 작은경우) 0을 리턴한다.

else-> poll한 작업을 1감소시켜 다시 PQ에 추가하고, 작업을 수행했으니 n도 1감소 시켜준다.

 

이렇게 반복을 돌게되면 결국 PQ의 특징때문에 제일 큰 값이 게속해서 poll된다.

그러므로, 할 수 있는 만큼 작업을 했을 때 남은 작업들의 제곱의 합을 최소화 할 수 있게 된다.