[Level 3] 5. 야근 지수 알고리즘

5. 야근 지수 알고리즘

Algorithm - Level 3

- 풀이

1. 원소 중 최대값에 1을 빼는 행동을 n번 만큼 반복하면 된다.

#include <string>
#include <vector>
#include <algorithm>
#include <numeric>
 
using namespace std;
 
long long calculate(vector<int> a){
    long long answer = 0;
    for(int i=0; i<a.size(); ++i)
        answer += (a[i]*a[i]);
    return answer;
}
 
long long solution(int n, vector<int> works) {
    long long answer = 0;
 
    sort(works.rbegin(), works.rend());
    
    if(works[0] == 0) return 0;
    else if(accumulate(works.begin(), works.end(), 0) < n) return 0;
    else if(works[0]-works[1]>=n){
        works[0]-=n;
        return calculate(works);
    } 
    else{
        for(int i=1; i<works.size(); ++i){
            while(works[i-1]!=works[i]){
                for(int j=0; j<i; ++j){
                    if(works[j]!=0){
                        works[j]--;
                        if(--n==0) return calculate(works);
                    }
                }
            }
        }
 
        if(works[0]>0){
            int idx=0;
            while(n>0){
                if(works[idx]>0){
                    works[idx++]--;
                    n--;
                }
            }
        }
    }
    
    return calculate(works);
}

위 풀이 방법은 <1>에서 말한 방법을 푼 방식으로 효율성을 위해 내림차순으로 정렬 후

1.  맨 처음 원소가 0이라면, 결과는 0이다.

2. 모든 원소의 덧셈이 n보다 이하라면 결과는 0이다.

3. 첫 원소와 두 번째 원소의 차이가 n 이상이라면 첫 원소에 n을 뺀 후 누적을 계산하면 된다.

4. 모든 원소를 같게한 후 남은 n만큼 번갈아가며 감소한다.

여기서 내가 원하는 함수(최대값을 찾는)를 못찾았었는데, 다른 사람의 풀이를 보니 내가 구현한 코드가 이미 함수로 구현되어 있었다. 그것을 이용하여 풀어보면,

#include <string>
#include <vector>
#include <iostream>
#include <algorithm>
using namespace std;
 
long long solution(int n, vector<int> works){
    long long answer = 0;
 
    while(n!=0){
        *max_element(works.begin(), works.end()) -= 1;
        n--;
    }
 
    for(int i=0; i<works.size(); ++i)
        answer += (works[i]*works[i]);
    
    return answer;
}

다음과 같이 간략화된다.

max_element는 최대값 원소를 반환하는 함수로, 여기에 *를 붙여 해당 위치에 1을 감산하여 반복할 수 있게 된다.