야근 지수

문제

회사원 Demi는 N시간 동안 야근 피로도를 최소화하도록 일하려고 합니다. 야근 피로도는 남은 일의 작업량을 제곱하여 더한 값입니다. 1시간 동안 작업량 1을 처리할 수 있을 때, 퇴근까지 남은 N 시간과 각 일에 대한 작업량 배열이 주어지면 야근 피로도의 최솟값을 구하는 프로그램을 작성하는 문제입니다.

• 입력: 퇴근까지 남은 시간 n (1,000,000 이하), 각 일에 대한 작업량 배열 works (길이 20,000 이하, 원소 50,000 이하)
• 출력: 야근 피로도를 최소화한 값(제곱의 합)을 반환합니다. 남은 작업량이 없다면 0을 반환합니다.

링크: https://school.programmers.co.kr/learn/courses/30/lessons/12927

풀이: 우선순위 큐(Max Heap)를 활용한 최댓값 감소 로직

어떤 수들의 제곱의 합을 최소화하려면, 수들 간의 편차를 줄여야 합니다. 즉, 가장 큰 값을 지속적으로 깎아내려 전체적인 수들을 평탄하게 만드는 것이 핵심입니다. 이를 위해 매시간마다 가장 큰 작업량을 찾아 1씩 줄이는 과정을 반복해야 하며, 최댓값을 효율적으로 찾고 관리하기 위해 우선순위 큐를 활용합니다.

• 사전 예외 처리: 먼저 works의 모든 작업량의 합을 구합니다. 이 총합이 남은 시간 n보다 작거나 같다면, 야근할 필요 없이 모든 일을 끝낼 수 있으므로 피로도는 0이 됩니다.
• Max Heap 초기화: 자바의 PriorityQueue는 기본적으로 오름차순(Min Heap) 정렬이므로, Collections.reverseOrder()를 적용하여 내림차순(Max Heap)으로 생성합니다. 배열의 모든 작업량을 큐에 넣습니다.
• 작업량 감소 (Greedy): 남은 시간 n만큼 반복문을 돌면서, 큐에서 가장 큰 값(poll)을 꺼냅니다. 꺼낸 값을 1 감소시킨 뒤 다시 큐에 넣습니다(add). 우선순위 큐의 특성상 항상 가장 큰 값이 루트에 위치하므로 매번 최댓값을 O(log M)의 시간 복잡도로 가져올 수 있습니다.
• 피로도 계산 및 오버플로우 주의: n시간이 모두 지나면 큐에 남은 요소들을 하나씩 꺼내어 제곱한 값을 누적 합산하여 반환합니다. 이때 작업량은 최대 50,000이 될 수 있으므로, 제곱을 할 경우 int형의 최댓값(약 21억)을 초과할 수 있습니다. 따라서 값을 꺼낼 때 long 타입으로 형변환을 한 뒤 제곱해야 오버플로우를 방지할 수 있습니다.

코드

import java.util.*;

class Solution {
    public long solution(int n, int[] works) {
        long sum = 0;
        for (int work : works) {
            sum += work;
        }
        
        // 모든 일을 처리할 수 있는 경우
        if (sum <= n) {
            return 0;
        }
        
        // 우선순위 큐 (내림차순 정렬, Max Heap)
        PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());
        for (int work : works) {
            pq.add(work);
        }
        
        // n시간 동안 가장 큰 작업량을 1씩 감소
        for (int i = 0; i < n; i++) {
            int x = pq.poll();
            pq.add(x - 1);
        }
        
        long ans = 0;
        // 남은 작업량들의 제곱의 합 계산
        while (!pq.isEmpty()) {
            // 오버플로우 방지를 위해 long으로 캐스팅 후 계산
            long x = pq.poll();
            ans += x * x;
        }
        
        return ans;
    }
}

알고리즘 분류

자료 구조 (Data Structures), 우선순위 큐 (Priority Queue), 탐욕법 (Greedy)