본문 바로가기
coding test

[programmers/level3] 야근 지수

by objet 2024. 1. 12.

 

문제 분석

야근 지수를 최소화하기 위해서는 각 요소 간 차이를 최소화해야 함. works 배열 안에 담겨 있는 요소들이 각각 작아지도록.

 

원래는 아래와 같은 방식으로 sort()를 이용해서 풀려고 했으나, 테스트 불통.

#include <string>
#include <vector>
#include <algorithm>
#include <iostream>

using namespace std;

long long solution(int n, vector<int> works) {
    long long answer = 0;
    
    // 각 works 안에 담겨 있는 요소들이 최대한 모두 작아지도록
    sort(works.rbegin(), works.rend());
    for (int i = 0; i < works.size(); i++)
        cout << works[i] << " ";
    cout << endl;
    
    if (works[0] - works[1] >= n)
    {
        // 첫번째한테 몰빵
        works[0] = works[0] - n;
    }
    else
    {
        // (n / works.size) 해서 몫의 값으로 각 원소 값 빼주기. 나머지값이 있으면 그 나머지값번째 요소까지 몫의 값을 한 번 더 더함.
        int stepsize = n / works.size();
        int terminal = n % works.size() - 1;
        
        if (stepsize == 0)
            for (int i = 0; i <= terminal; i++)
                works[i] -= 1;
        else 
        {
            for (int i = 0; i < works.size(); i++)
            {
                works[i] -= stepsize;
                if (i <= terminal)
                    works[i] -= stepsize;
                if (works[i] < 0)
                    works[i] = 0;
            }
        }
    }
    
    for (int i = 0; i < works.size(); i++) {
        cout << works[i] << " ";
        answer += works[i] * works[i];
    }
    
    return answer;
}

 

이후 우선순위 큐에 대해 공부하고, 우선순위 큐로 풀면 쉽겠다는 깨달음을 얻음.

내림차순으로 우선순위를 준 다음 그때그때 가장 큰 수를 top()으로 꺼내서 1씩 줄여주면 된다.

#include <string>
#include <vector>
#include <queue>
#include <algorithm>
#include <iostream>

using namespace std;

long long solution(int n, vector<int> works) {
    long long answer = 0;
    
    // 각 works 안에 담겨 있는 요소들이 최대한 모두 작아지도록
    priority_queue<int> pq(works.begin(), works.end());
    while (n-- && pq.top() > 0)
    {
        int top = pq.top();
        pq.pop();
        pq.push(top - 1);
    }
    
    while (!pq.empty()) 
    {
        answer += pq.top() * pq.top();
        pq.pop();
    }
    
    return answer;
}

 

위처럼 풀었더니 통과!