Solution.java

class Solution {
    public long solution(int n, int[] works) {
        long answer = 0;
        
        int[] tmp = new int[50001];
        
        for (int work : works) {
            tmp[work]++;
        }
        
        int i;
        for (i = 50000; i > 0; i--) {
            if (tmp[i] < n) {
                n -= tmp[i];
                tmp[i - 1] += tmp[i];
            }
            else {
                tmp[i] -= n;
                tmp[i - 1] += n;
                break;
            }
        }
        
        while (i > 0) {
            answer += (long) i * i * tmp[i];
            i--;
        }
        
        return answer;
    }
}

출처: 프로그래머스 코딩 테스트 연습, https://programmers.co.kr/learn/challenges