K smallest elements in an array

Find K smallest elements or largest elements in a set or array is very frequently asked interview question in Amazon and Microsoft interviews. This question is interesting because interviewer asks to optimize the solution in terms of memory and time complexity.

Question: Given an array of N numbers which are unsorted, find K smallest elements in that array. For example if the array A = [3,4,8,6,5,1,2,3,5], and if K =3;  then K smallest elements in this array are 1,2 and 3.

K smallest elements using sorting

What if input array is sorted?  Then, numbers from index 0 to K-1 are K smallest elements in array. If array is unsorted, we first sort array and then return first K elements from it.

With best sorting algorithm,  time complexity of sorting is O(n log n) and then to return K elements requires O(K) time. Hence, total complexity of method is O(n log n).

K smallest elements  : Sorting base implementation

To see the complete implementation of quick sort please read this post : QuickSort Algorithm Can we optimize it?

K smallest elements using partial quicksort

Think of this: Do we want sort entire array before, K smallest elements can be found out? The answer is NO. In quicksort, we find correct position of pivot. Pivot is placed at a position in array such that all elements on to left of pivot are smaller and all elements on to right are greater than pivot. So, if pivot is placed at jth index of array, then we already know j smallest elements in array, from index 0 to j. Now we just need to check if j is equal to K. If yes, don’t sort array further. Just return elements from 0 to index of pivot. In worst case, we will end up sorting the entire array but in most cases, we will save computing cycles. Implementation for this method is as given below.

K smallest elements using heaps

Can we optimize it further? Whenever there problem is of K smallest elements, think in terms of heaps. Before reading further, I would suggest to read heaps basics here if you are not very comfortable with heaps.

To solve this particular problem, we will create a heap with N elements of array. Which heap should we build? Since we need K minimum elements, numbers from heap should be taken out in increasing order. Which heap gives us numbers in increasing order when root of heap is taken out one by one?  Yes, it is min heap. Hence, we create a min heap with N elements.

Now we take out minimum element from the heap. After removing root, we will heapify min heap again, so that next minimum in the heap becomes the root. We will repeat the process for K times and we will have K smallest elements with us.

Question which is usually asked by interviewer is: what is complexity of this method? In order to answer that let see how many steps we have performed. First step is to create a heap from N elements array, whose complexity is O(N). Usually people have confusion that to build a heap from N elements array is O(nlogn) which is not true. To further understand why that is so, please follow the post: heaps basics

Next step is to take out the root K times. Each time heapify is called on heap, which has complexity of O(log n), so each removal costs us (log n) and K removal will have complexity of O(K log n).

Hence overall complexity will be O(n) + O(k log n). If n is very large as compared to K then complexity will be O(N).

Problem with this method to find K smallest elements is that we need to maintain a heap of N elements. If the number of elements in array are in tunes of millions and if just 1o smallest elements are to be found, it will be a cumbersome process to maintain a heap of millions of nodes and huge waste of memory.

K smallest elements using heaps but with a tweak

Since we are interested in only K elements, why don’t we make a heap of K smallest elements we have seen so far?

Take first K elements and then create a heap. What heap should we build now? We need to created a max heap here, because we will scan entire array and at each element of array, we will check if it is less than the max on heap. If yes, then that element needs to be inserted into heap and maximum element from heap is deleted. Once scan is complete, heap will have K smallest elements.

Well what is complexity of this method? To create a heap with K elements it will take O(K) time. Each removal of root and insertion a of new element to heap has complexity of O(log K). Over all complexity becomes O(N log K) + O(K). This is better when we have use set of integers and there is not enough memory to bring them all in at once as we did in last method.

K smallest elements Implementation

Bonus question: Given a set of points (x,y) in two dimensional plane, find K nearest points from the origin? Can you map this problem to above problem? If yes, please send me your solution and I will put it here.

Same principle of finding k smallest elements in an array can be applied here too. Try it and comment your solution if you need feedback.