educative.io

Max heap alternative

Courses
,
#1

Hi all,
I learned this technique from a mock interview with my friend.
instead of min heap, we can use a max heap and adjust the priority queue to only contain k elements.
then your kth smallest element at the end is just the first one to come out from the priority queue.

public static int findKthSmallest(List<Integer[]> lists, int k) {
    Queue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());

    for (Integer[] arr : lists) {
      for (int i = 0; i < arr.length; ++i) {
        pq.offer(arr[i]);

        if (pq.size() > k) {
           pq.poll();
        }
      }
    }

    return pq.poll();
}
#2

example [[2, 6, 8], [3, 6, 7], [1, 3, 4]], k = 5
correct answer = 4

each iteration

  1. add 2, size = 1, [2]
  2. add 6, size = 2, [6, 2]
  3. add 8, size = 3, [8, 6, 2]
  4. add 3, size = 4, [8, 6, 3, 2]
  5. add 6, size = 5, [8, 6, 6, 3, 2]
  6. add 7, size = 6, [8, 7, 6, 6, 3, 2]
    ^size is larger than k, we poll the largest off and becomes [7, 6, 6, 3, 2] size = 5
  7. add 1, size = 6, [7, 6, 6, 3, 2, 1]
    ^same story, poll 1 largest and becomes [6, 6, 3, 2, 1] size = 5
  8. add 3, size = 6, [6, 6, 3, 3, 2, 1]
    ^same story, poll 1 largest and becomes [6, 3, 3, 2, 1] size = 5
  9. add 4, size = 6, [6, 4, 3, 3, 2, 1]
    ^same story, poll 1 largest and becomes [4, 3, 3, 2, 1] size = 5

The next poll is the 5th largest in these arrays which is 4.
This idea is to maintain the pq with size k. we add one and toss one away that we dont need.
TC: O(N log K) <= we take N times adding into a k sized PQ
SC: O(K) <= since we control the PQ size to k

#3

This is a solution for when the lists are not ordered, right? When the lists are ordered you can finish as soon as the Kth value is found.