Time complexity of Kth smallest using Heap

Question

Following is a code for finding the kth smallest element in an array using heap. The time complexity is O(n log(k)) where k is the size of the heap.

As per my understanding, you first go through the entire array i.e. O(n) to populate your heap. And, when you reach the end of the array, you would have the kth smallest element at the top of the heap that you can instantly return as the final answer.

However, in the code below, there is an additional loop starting from k to the length of the array. I'm not really understanding the need for this second loop.

public int findKthSmallest(int[] arr, int k ) {

    if(k <= 0 || k > arr.length) {
        throw new IllegalArgumentException();
    }

    PriorityQueue<Integer> smallestK = new PriorityQueue<>(k, Collections.reverseOrder());

    for(int i = 0; i < arr.length; i++) {
        smallestK.add(arr[i]);
    }

    for(int j = k; j < arr.length; j++) {
        if(arr[j] < smallestK.peek()) {
            smallestK.remove();
            smallestK.add(arr[j]);
        }
    }
    return smallestK.peek();
}

Answer 1

You read the wrong code, it should be:

for(int i = 0; i < k; i++) {
    smallestK.add(arr[i]);
}

In the first loop we need to insert the first k element in heap.

At current moment, smallestK.peek() will hold the current smallestK.

In the second loop, we process the remaining element in array. We compare the value with the current smallestK.