如何在多維數組中找到k個最小數的索引？

Question

我想形成包含k個最小值的指數在一個數組的數組：

import heapq 
import numpy as np 

a= np.array([[1, 3, 5, 2, 3], 
     [7, 6, 5, 2, 4], 
     [2, 0, 5, 6, 4]]) 

[t[0] for t in heapq.nsmallest(2,enumerate(a[1]),lambda(t):t[1])] 
===[3, 4] 

但這種失敗：

[t[0] for t in heapq.nsmallest(2,enumerate(a.all()),lambda(t):t[1])] 

Traceback (most recent call last): 
    File "<pyshell#19>", line 1, in <module> 
    [t[0] for t in heapq.nsmallest(2,enumerate(a.all()),lambda(t):t[1])] 
TypeError: 'numpy.bool_' object is not iterable 

Answer 1

可以使用numpy.ndenumerate用堆或者部分排序的建議由大衛：

a = np.array([[1, 3, 5, 2, 3], 
     [7, 6, 5, 2, 4], 
     [2, 0, 5, 6, 4]]) 
heap = [(v, k) for k,v in numpy.ndenumerate(npa)] 
heapq.heapify(heap) 
heapq.nsmallest(10, heap) # for k = 10 

，你會得到：

[(0, (2, 1)), 
(1, (0, 0)), 
(2, (0, 3)), 
(2, (1, 3)), 
(2, (2, 0)), 
(3, (0, 1)), 
(3, (0, 4)), 
(4, (1, 4)), 
(4, (2, 4)), 
(5, (0, 2))] 

Answer 2

您的問題是在a.all()：

[t[0] for t in heapq.nsmallest(2,enumerate(a.all()),lambda(t):t[1])] 

這檢查數組中所有元素的真實性，即False（因爲您有0）。

如果數組與k相比不太大，則可以使用.argsort來獲取值。在這裏，我將選擇最大的兩個位置的每一行：

print a.argsort()[:,:2] 

array([[0, 3], 
     [3, 4], 
     [1, 0]]) 

如果你想在全球最低的位置，壓平陣列firts：

a.flatten().argsort()[:2] 

如果陣列是非常大的，您可以使用np.argpartition獲得更好的性能，這將僅執行部分排序。