[프로그래머스] 야근 지수

프로그래머스

코드 중심의 개발자 채용. 스택 기반의 포지션 매칭. 프로그래머스의 개발자 맞춤형 프로필을 등록하고, 나와 기술 궁합이 잘 맞는 기업들을 매칭 받으세요.

programmers.co.kr

풀이

직관적으로 생각하면 퇴근할 때까지 1시간마다 가장 많이 남은 일을 처리하면 된다. 1시간마다 배열을 정렬해서 가장 큰 값을 1 빼주는 방법을 떠올릴 수도 있는데, 이는 효율적이지 못하다. 이 문제에서 필요한 것은 정렬된 배열이 아니고 최댓값이므로 우선 순위 큐를 사용하면 훨씬 더 효율적으로 해결할 수 있다.

코드

C++

#include <string>
#include <vector>
#include <queue>
using namespace std;

long long solution(int n, vector<int> works) {
    long long answer = 0;
    priority_queue<int> pq;
    for (int i : works) pq.push(i);
    while (n--) {
        int tmp = pq.top();
        if (tmp > 0) {
            pq.pop();
            pq.push(tmp - 1);
        }
    }
    while (!pq.empty()) {
        int tmp = pq.top();
        pq.pop();
        answer += (tmp * tmp);
    }
    return answer;
}

Java

import java.util.Comparator;
import java.util.PriorityQueue;

class Solution {
    public long solution(int n, int[] works) {
        long answer = 0;
        PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.reverseOrder());
        for (int i : works) pq.add(i);
        while (n-- > 0) {
            if (pq.peek() > 0) {
                pq.add(pq.poll() - 1);
            }
        }
        while (!pq.isEmpty()) {
            int num = pq.poll();
            answer += (num * num);
        }
        return answer;
    }
}