[Algorithm] 야근 지수

문제

회사원 Demi는 가끔은 야근을 하는데요, 야근을 하면 야근 피로도가 쌓입니다. 야근 피로도는 야근을 시작한 시점에서 남은 일의 작업량을 제곱하여 더한 값입니다. Demi는 N시간 동안 야근 피로도를 최소화하도록 일할 겁니다.Demi가 1시간 동안 작업량 1만큼을 처리할 수 있다고 할 때, 퇴근까지 남은 N 시간과 각 일에 대한 작업량 works에 대해 야근 피로도를 최소화한 값을 리턴하는 함수 solution을 완성해주세요.

 

제한 사항

• works는 길이 1 이상, 20,000 이하인 배열입니다.
• works의 원소는 50000 이하인 자연수입니다.
• n은 1,000,000 이하인 자연수입니다.

 

입출력 예

[4, 3, 3]	4	12
[2, 1, 2]	1	6
[1,1]	3	0

 

 

풀이

처음 시도했던 접근 방향은 평균값을 내며 그리디처럼 푸는 방법이었다. 남은 작업량의 제곱값으로 피로도가 결정되기 때문에 각 남은 작업량의 차이가 완만하게 되야한다고 생각했다. 먼저 모든 작업량에서 n을 빼준 뒤에 for문으로 돌면서 평균값을 구하고 이를 반복문의 현재 값과 비교해서 answer에 더해주는 방법으로 시도했다. 예제는 맞았지만 제출을 하고보니 대부분의 케이스에서 실패를 했는데 이유를 생각해보니 현재값이 굉장히 크고 뒤의 값들이 대부분 작은 경우 현재 값에서 많이 빼줘야하지만 평균 값에 영향을 받아 요구하는 값만큼의 처리를 못하는 것을 알게 되었다. 이런 시도 이후 그렇다면 가장 큰 값에서부터 하나씩 빼줘야겠다라고 생각했다.

 

두 번째 시도로는 while문으로 n이 0 초과일때 돌아가도록 하고 남은 작업량들을 내림차순으로 정렬해서 제일 앞에 있는 값에서 빼주는 방법이었다. 하지만 매번 sort 메서드를 호출한다면 시간초과로 이어질 것이라고 생각했고 구현했었던 max heap을 사용했다. while문이 종료되면 더 이상 뺄 값이 없거나 빼지 않아도 되는 상황이므로 heap에 존재하는 값을 answer에 제곱해서 넣어줌으로서 통과할 수 있었다.

 

코드

class MaxHeap {
  constructor() {
    this.heap = [];
    this.heap.push(1e9);
  }

  insert(value) {
    this.heap.push(value);
    this.upheap(this.heap.length - 1);
  }

  upheap(pos) {
    let temp = this.heap[pos];

    while (temp > this.heap[parseInt(pos / 2)]) {
      this.heap[pos] = this.heap[parseInt(pos / 2)];
      pos = parseInt(pos / 2);
    }
    this.heap[pos] = temp;
  }

  get() {
    if (this.size() === 0) return false;
    if (this.size() === 1) return this.heap.pop();
    const result = this.heap[1];
    this.heap[1] = this.heap.pop();
    this.downheap(1, this.heap.length - 1);
    return result;
  }

  downheap(pos, len) {
    let temp = this.heap[pos],
      child;
    while (pos <= parseInt(len / 2)) {
      child = pos * 2;
      if (child < len && this.heap[child] < this.heap[child + 1]) child++;
      if (temp >= this.heap[child]) break;
      this.heap[pos] = this.heap[child];
      pos = child;
    }
    this.heap[pos] = temp;
  }

  size() {
    return this.heap.length - 1;
  }
}

function solution(n, works) {
  var answer = 0;
  const len = works.length;

  const heap = new MaxHeap();
  for (const work of works) heap.insert(work);

  while (n > 0) {
    let num = heap.get();
    if (!num) return 0;

    heap.insert(num - 1);
    n -= 1;
  }
  for (let i = 1; i < heap.size() + 1; i++) answer += heap.heap[i] ** 2;

  return answer;
}