Code Yarns ‍👨‍💻
Tech BlogPersonal Blog

How to find indices of smallest K items

📅 2016-Jun-16 ⬩ ✍️ Ashwin Nanjappa ⬩ 🏷️ partial_sort, priority queue, puzzle ⬩ 📚 Archive


Recently I actually faced a problem, the likes of which you see being asked in interviews. Input is a array of N values that is not sorted. We want this array to be untouched, since it is required in the given order for later use. What we want to find out is the indices of the smallest K values in this unsorted array. The K in this case is very small compared to N.

Solution 1

The straightforward solution is:

The computational complexity of this solution can probably not be beaten. However, the problems with this solution for my particular application were:

Solution 2

As you can imagine, the ideal solution would have complexity similar to the above solution. But practically it should not require N new storage, but K size storage. In fact, this can be solved using a std::priority_queue. Here is the pseudocode of my solution:

GetSmallestKIndices(input_array, K):
    Create priority queue pq

    Insert first K pairs of items and indices into pq
    (pq will be maintained at size K)

    Iterate over each of the N-K items in input_array
        If item is bigger than top item in pq then

        We have now found item smaller than items in pq
        Pop off the top item from pq
        Insert pair of current item and its index

    What we have left in pq is smallest K values and their indices

It goes without saying that this solution is only ideal for this particular setup where K is very small compared to N.

© 2022 Ashwin Nanjappa • All writing under CC BY-SA license • 🐘📧