[Programmers] 야근 지수

프로그래머스 - 야근 지수 문제입니다.

 

문제

각 work의 제곱을 더한 값이 최소가 되려면 가장 큰 works 부터 줄여야겠다는 생각을 했습니다.

max heap을 이용하여 n과 works를 1씩 줄이면 답을 낼 수 있고, 효율성을 올리려면 가장 큰 work 와 두번째로 큰 work의 차이를 한번에 n에서 빼면 될 것 같습니다.

 

#include <vector>
#include <queue>

using namespace std;

long long solution(int n, vector<int> works) {
    priority_queue<int> pq;		// max heap
    long long answer = 0;
    int preValue;
    
    for(int i = 0; i < works.size(); i++) {
        pq.push(works[i]);
    }

    while(!pq.empty() && n > 0) {	// 큐가 비어있지 않고 n이 0보다 클 때
        preValue = pq.top();
        pq.pop();
        
        n--;
        preValue--;
        if(preValue != 0) pq.push(preValue);	// n과 work를 1씩 뺀 후 push
    }
    
    if(pq.empty()) answer = 0;	// 큐가 비었다면 남은 work가 없으므로 0
    else {
        while(!pq.empty()) {
            answer += (pq.top() * pq.top());
            pq.pop();
        }
    }
    return answer;
}