プログラマー-残業指数

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

問題の説明

会社員のデミはたまに残業しますが、残業すると残業疲労度が溜まります.残業疲労度は、残業開始時の残量に1つの数値を乗じたものです.デミはN時間以内に残業疲労度を最小限に抑える.デミが1時間以内に1時間の仕事量を処理できると言った場合、退勤前の残りのN時間と各仕事の仕事量について、残業疲労度を最低値に下げる関数解を返してください.
せいげんじょうけん
worksは、長さが1より大きく、20000未満の配列です.
worksの要素は50000以下の自然数です.
nは1000000以下の自然数である.
入力例
works = [4, 3, 3] n = 4
works =[2, 1, 2] n = 1
works =[1,1] n = 3
サンプル出力
12
6
0
I/O例説明
I/O例#1
n=4の場合、残りの仕事量が[4,3]の場合、残業指数を最小化するために4時間働いた結果は[2,2,2]であった.このとき残業指数は22+22+22=22です.
I/O例#2
n=1の場合、残りの仕事量が[2,1,2]であれば、残業指数を最小化するために1時間働いた結果は[1,1,2]である.残業指数は12+12+22=6です.

能率の問題を解決する

import java.util.*;

class Solution {
    public long solution(int n, int[] works) {
        long answer = 0;
        int total = Arrays.stream(works).sum();
        if(total <= n)
            return 0;
        Arrays.sort(works);
        int index = works.length - 1;
        while(index > 0 && n > 0) {
            while(works[index] > works[index - 1] && n > 0) {
                for(int i = works.length - 1; i >= index && n > 0; i--) {
                    works[i]--;
                    n--;
                }
            }
            index--;
        }
        while(n > 0) {
            for(int i = 0; i < works.length; i++) {
                works[i]--;
                n--;
                if(n == 0)
                    break;
            }
        }
        return Arrays.stream(works).map(i -> i = (int) Math.pow(i, 2)).sum();
    }
}

最大から縮小までが近づいてきましたが、優先順位Queueとは思いませんでしたので、このようにパターンに近づきました.
ex)
1 3 4 5 7時
9->7=13 4 5 7に減少
7->5=1 3 4 5 5に減少
.
.
これで最後に1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
解題方法は正しいが、タイムアウトした.

priorityQueueの使用ヒントを聞いてから再解釈

import java.util.*;

class Solution {
    public long solution(int n, int[] works) {
        long answer = 0;
        PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());
        
        Arrays.stream(works).forEach(i -> pq.add(i));
        while(!pq.isEmpty() && n > 0) {
            int temp = pq.poll();
            temp -= 1;
            if(temp > 0)
                pq.add(temp);
            n--;
        }
        
        if(n > 0)
            return answer;
        else {
            while(!pq.isEmpty())
                answer += (int) Math.pow(pq.poll(), 2);
            return answer;
        }
    }
}