프로그래머스 - 12927번: 야근 지수 [Java]

문제: https://school.programmers.co.kr/learn/courses/30/lessons/12927

프로그래머스

1. 코드

import java.util.*;
class Solution {
    public long solution(int n, int[] works) {
        long answer = 0;
        PriorityQueue <Integer> queue = new PriorityQueue<>(Collections.reverseOrder());
    	for(int w : works)
            queue.add(w);
        for(int i = 0; i < n; i++)
            queue.offer(queue.poll() - 1);
        
        for(int q : queue)
            if(q > 0)
                answer += q * q;
        return answer;
    }
}

2. 설명

n만큼의 수를 works에 균등하게 나누어야 하는데 이 방식을 힙 정렬을 이용하여 해결할 수 있습니다.

long answer = 0;
PriorityQueue <Integer> queue = new PriorityQueue<>(Collections.reverseOrder());
for(int w : works)
    queue.add(w);

최대 힙을 기준으로 우선순위 큐를 생성합니다. 그리고 반복문을 통해서 works의 값들을 우선순위 큐에 저장합니다.

for(int i = 0; i < n; i++)
    queue.offer(queue.poll() - 1);

for(int q : queue)
    if(q > 0)
        answer += q * q;

n만큼 우선순위 큐에서 최대값 꺼내서 -1을 한 뒤 다시 우선순위에 삽입합니다.

만약 0미만의 수, 즉 -1이하의 수가 제곱이 되면 양수가 되어 answer에 영향을 미치기에 q의 값이 양수일 경우에만 제곱한 값을 answer 더한 뒤 answer를 반환하면 쉽게 해결됩니다.

3. 정리

1. 힙 정렬을 하여 시간초과를 해결할 수 있다.
2. 우선순위 큐를 적절히 활용하자.

출처: 프로그래머스 코딩 테스트 연습, 
https://school.programmers.co.kr/learn/challenges