[우선순위 큐] 703. Kth Largest Element in a Stream

문제 바로가기

 

정수로 이루어진 스트림에서 k번째로 큰 숫자를 찾는 문제입니다. KthLargest 클래스의 생성자와,  add() 메서드를 구현하면 됩니다.

 

class KthLargest {
  private PriorityQueue<Integer> minHeap;
  int k;

  public KthLargest(int k, int[] nums) {
    minHeap = new PriorityQueue<>();
    this.k = k;
    for (int num : nums) {
      add(num);
    }
  }

  public int add(int val) {
    if (minHeap.size() < k) {
      minHeap.add(val);
    } else {
      if (minHeap.peek() < val) {
        minHeap.poll();
        minHeap.add(val);
      }
    }

    return minHeap.peek();
  }
}

 

이 문제의 핵심은 우선순위 큐(최소 힙)을 활용하는 것입니다. add() 메서드가 호출될 때마다 매개변수로 들어오는 정수도 스트림에 같이 넣어주면서 항상 k번째로 큰 수를 반환해야 하는데요, add() 메서드가 호출될 때마다 정렬을 한다면 아마 시간복잡도에서 걸릴겁니다.

 

최소 힙을 어떻게 써먹냐~ k번째로 큰 수를 찾아야 하는데 최대 힙을 써야 하는게 아니냐~ 싶을 수 있지만.. 반대로 생각해야 합니다. 바로 최소 힙의 원소를 항상 k개로 유지만 해준다면 최소 힙의 가장 위에 있는 원소가 바로 우리가 찾는 k번째로 큰 수가 되는 것이죵.

 

add() 메서드만 살펴보겠습니다.

최소 힙의 크기가 k보다 작으면 당연히 최소 힙에 넣어주고, 그렇지 않다면?! peek() 과 새로 들어온 숫자를 비교해줍니다. 만약 새로 들어온 숫자가 더 작다면 얘는 k + 1번째로 큰 수가 되겠죠. 반대라면 k번째로 큰 수가 될테고, 원래 있던 k번째 큰 수를 쫓아내주면 됩니당.

 

 

 

감사합니다!!!!~!~!