Python algorithm to find the indexes of the k smallest number in an unsorted array?

Question

Is there any algorithm to find the indexes of the k smallest numbers in an unsorted array in python? I know how this can be achieved using numpy module, but I am not looking for that. One direction that immediately comes to my mind is that it has to with sorting algorithms. So lets say that I have an algo to sort an array in python using bubble sort:

def bubbleSort(arr):
n = len(arr)

# Traverse through all array elements
for i in range(n):

    for j in range(0, n-i-1):
        # Swap if the element found is greater
        # than the next element
        if arr[j] > arr[j+1] :
            arr[j], arr[j+1] = arr[j+1], arr[j]

I am not sure how to modify this algorithm to just return the indexes of the k smallest number in the array. Any help using sorting algorithm or selecting algorithm like quickselect, quicksort shall be appreciated.

EDIT 1 : So say array is:

a = [12, 11, 0, 35, 16, 17, 23, 21, 5]

Then it must just return an array: index_of_least_k = [2,8,1]

for k = 3.

If I had to modify the sorting algorithm say bubble sort, I understand how to change so that it swaps the indexes this time, say:

def modified_bubbleSort(arr, index):
      n = len(arr)

      # Traverse through all array elements
      for i in range(n):

           for j in range(0, n-i-1):
                  # Swap if the element found is greater
                  # than the next element
                  if arr[j] > arr[j+1] :
                         index[j], index[j+1] = index[j+1], index[j]
      return index


array = [12, 11, 0, 35, 16, 17, 23, 21, 5]
index = [0, 1, 2, 3, 4, 5, 6, 7, 8]

indexOfAllsorted = modified_bubblesort(array, index)

In this case it returns me:

indexOfAllsorted = [2,8,1,0,4,5,7,6]

I do not want that because there are extra 5 values, to avoid the memory overhead my algorithm should just have:

index_of_least_k = [0, 0, 0]

in the memory for k =3 and then fill it as it proceeds. I hope I made it clear.

EDIT2: I am not looking for any library or modules to accomplish that in python.

Answer 1

You can use heapq.nsmallest to get the n smallest items from an iterable. So how do you create an iterable such that it measures the values of the input, but returns their indices? One way is to use the enumerate function to get an iterable of (index, value) pairs, then use a key function to only use the values.

from heapq import nsmallest
from operator import itemgetter

def indices_of_n_smallest(n, seq):
    smallest_with_indices = nsmallest(n, enumerate(seq), key=itemgetter(1))
    return [i for i, x in smallest_with_indices]

array = [12, 11, 0, 35, 16, 17, 23, 21, 5]
indices_of_n_smallest(3, array)
# [2, 8, 1]

Answer 2

Here's the thing about bubble sort. Each time the inner loop finishes iterating, exactly one element finds its correct position. Your code for example, finds the ith largest element each time, as it's sorting in ascending order. Let's flip that > sign to <; now it'll find the ith smallest element every time the j loop finishes. So if you stop the sorting when i=k, you'll have your k smallest elements.

def modified_bubbleSort(arr, index, k):
  n = len(arr)
  ans = []

  for i in range(k):

       for j in range(0, n-i-1):
              # Swap if the element found is smaller
              # than the next element
              if arr[index[j]] < arr[index[j+1]] :
                     index[j], index[j+1] = index[j+1], index[j]
       ans.append(index[n-i-1])
  return ans