Python算法在未排序的数组中找到k个最小数字的索引？

Question

有没有算法可以在python中找到未排序数组中k个最小数字的索引？我知道如何使用 numpy 模块实现这一点，但我不是在寻找它。我立即想到的一个方向是它必须与排序算法有关。所以假设我有一个算法可以使用冒泡排序在 python 中对数组进行排序:

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]

我不确定如何修改此算法以仅返回数组中 k 最小数字的索引。任何使用排序算法或选择算法(如 quickselect、quicksort)的帮助将不胜感激。

编辑 1:所以说数组是:

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

然后它必须只返回一个数组: index_of_least_k = [2,8,1]

对于 k = 3。

如果我不得不修改排序算法，比如冒泡排序，我知道如何更改以便这次交换索引，比如:

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)

在这种情况下它返回我:

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

我不想要那个，因为有额外的 5 个值，为了避免我的算法应该有的内存开销:

index_of_least_k = [0, 0, 0]

在内存中 k =3，然后在进行时填充它。我希望我说清楚了。

EDIT2:我不是在寻找任何库或模块来在 python 中完成它。

Answer 1

您可以使用 heapq.nsmallest从可迭代对象中获取 n 最小项。那么如何创建一个可迭代对象来测量输入的值，但返回它们的索引呢？一种方法是使用 enumerate 函数获取 (index, value) 对的可迭代对象，然后使用键函数仅使用值。

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]