LeetCode/java/kth-largest-element-in-a-stream.java

import java.util.PriorityQueue;

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

  public KthLargest(int k, int[] nums) {
    this.k = k;
    heap = new PriorityQueue<>(k + 1);

    for (var num : nums) {
      add(num);
    }
  }

  public int add(int val) {
    if (heap.size() < k || heap.peek() < val) {
      heap.add(val);

      if (heap.size() > k) {
        heap.remove();
      }
    }

    return heap.peek();
  }
}