在 numpy 数组中找到第 n 个最小的元素

新手上路,请多包涵

我需要在 1D numpy.array 中找到最小的第 n 个元素。

例如:

 a = np.array([90,10,30,40,80,70,20,50,60,0])

我想获得第 5 个最小的元素,所以我想要的输出是 40

我目前的解决方案是这样的:

 result = np.max(np.partition(a, 5)[:5])

然而,找到 5 个最小的元素然后取其中最大的一个对我来说似乎并不笨拙。有更好的方法吗?我错过了一个可以实现我的目标的功能吗?

有些问题的标题与此类似,但我没有看到任何可以回答我的问题的问题。

编辑:

我本来应该提到它,但性能对我来说非常重要;因此, heapq 解决方案虽然不错,但对我来说不起作用。

 import numpy as np
import heapq

def find_nth_smallest_old_way(a, n):
    return np.max(np.partition(a, n)[:n])

# Solution suggested by Jaime and HYRY
def find_nth_smallest_proper_way(a, n):
    return np.partition(a, n-1)[n-1]

def find_nth_smallest_heapq(a, n):
    return heapq.nsmallest(n, a)[-1]
#
n_iterations = 10000

a = np.arange(1000)
np.random.shuffle(a)

t1 = timeit('find_nth_smallest_old_way(a, 100)', 'from __main__ import find_nth_smallest_old_way, a', number = n_iterations)
print 'time taken using partition old_way: {}'.format(t1)
t2 = timeit('find_nth_smallest_proper_way(a, 100)', 'from __main__ import find_nth_smallest_proper_way, a', number = n_iterations)
print 'time taken using partition proper way: {}'.format(t2)
t3 = timeit('find_nth_smallest_heapq(a, 100)', 'from __main__ import find_nth_smallest_heapq, a', number = n_iterations)
print 'time taken using heapq : {}'.format(t3)

结果:

 time taken using partition old_way: 0.255564928055
time taken using partition proper way: 0.129678010941
time taken using heapq : 7.81094002724

原文由 Akavall 发布,翻译遵循 CC BY-SA 4.0 许可协议

阅读 531
2 个回答

除非我遗漏了什么,否则你要做的是:

 >>> a = np.array([90,10,30,40,80,70,20,50,60,0])
>>> np.partition(a, 4)[4]
40

np.partition(a, k) will place the k+1 -th smallest element of a at a[k] , smaller values in a[:k] and larger a[k+1:] 中的值。唯一需要注意的是,由于 0 索引,第五个元素位于索引 4 处。

原文由 Jaime 发布,翻译遵循 CC BY-SA 4.0 许可协议

您可以使用 heapq.nsmallest

 >>> import numpy as np
>>> import heapq
>>>
>>> a = np.array([90,10,30,40,80,70,20,50,60,0])
>>> heapq.nsmallest(5, a)[-1]
40

原文由 arshajii 发布,翻译遵循 CC BY-SA 3.0 许可协议

撰写回答
你尚未登录,登录后可以
  • 和开发者交流问题的细节
  • 关注并接收问题和回答的更新提醒
  • 参与内容的编辑和改进,让解决方法与时俱进
推荐问题