프로그래머스/lv3/12927. 야근 지수

#include <string>
#include <vector>
#include <queue>
#include <numeric>
#include <math.h>
using namespace std;

long long solution(int n, vector<int> works) {
    long long answer = 0;
    priority_queue<int> q(works.begin(), works.end());
    int k = accumulate(works.begin(), works.end(), 1);

    while (n) {
        int z = q.top();
        q.pop();
        while (z >= q.top()) {
            z--, n--, k--;
            if (!k) return 0;
            if (!n) break;
        }
        q.push(z);
    }

    while (!q.empty()) {
        int i = q.top();
        q.pop();
        answer += pow(i, 2);
    }

    return answer;
}

평균값을 최대한 줄여야 제곱을 더하는 과정에서 가장 작은 값을 만들 수 있다. 그러므로 우선순위 큐를 사용해서 works의 모든 값을 할당해주고, k에 works의 모든 값을 더해 야근이 필요없을 때를 예외처리하기 위해 할당했다.

퇴근까지 남은 시간이 모두 소진될 때 까지 while문을 반복한다. 제일 큰 작업량인 q.top()을 z에 할당하고 pop하면 두번째로 큰 작업량이 q.top()보다 작아질 때 까지 z를 처리한다. 이 때 마찬가지로 예외처리를 하기 위해 k도 하나씩 줄이고, 남은 작업시간인 n도 똑같이 1씩 줄여나가야 한다.
처리가 완료된 z는 다시 작업해야 할 q에 다시 우선순위 큐에 push해주며 위 과정을 반복한다. 이렇게 되면 모든 값에 대해 평균값을 균일하게 줄여나갈 수 있다. 그 후 q의 모든 값을 제곱해 answer에 더해주고 return한다.