我想形成一个数组，其中包含一个数组中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
您可以使用numpy.ndenumerate堆，也可以使用David建议的部分排序：
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))]
本文收集自互联网，转载请注明来源。
如有侵权，请联系[email protected] 删除。
我来说两句