python - 找到与给定数字最接近的 k 个数字

Question

假设我有一个列表[1,2,3,4,5,6,7]。我想找到最接近 6.5 的 3 个数字。那么返回值将是 [5,6,7].

在 python 中找到一个最接近的数字并不是那么棘手，可以使用

min(myList, key=lambda x:abs(x-myNumber))

但我试图不围绕这个循环来查找 k 个最接近的数字。有没有pythonic的方式来完成上述任务？

Answer 1

简短的回答

heapq.nsmallest()函数将巧妙而有效地做到这一点:

>>> from heapq import nsmallest
>>> s = [1,2,3,4,5,6,7]
>>> nsmallest(3, s, key=lambda x: abs(x - 6.5))
[6, 7, 5]

基本上这就是说，“给我三个与数字 6.5 绝对差值最小的输入值”。

针对重复查找进行优化

@Phylliida 在评论中询问如何优化具有不同起点的重复查找。一种方法是对数据进行预排序，然后使用 bisect定位小搜索段的中心:

from bisect import bisect

def k_nearest(k, center, sorted_data):
    'Return *k* members of *sorted_data* nearest to *center*'
    i = bisect(sorted_data, center)
    segment = sorted_data[max(i-k, 0) : i+k]
    return nsmallest(k, segment, key=lambda x: abs(x - center))

例如:

>>> s.sort()
>>> k_nearest(3, 6.5, s)
[6, 7, 5]
>>> k_nearest(3, 0.5, s)
[1, 2, 3]
>>> k_nearest(3, 4.5, s)    
[4, 5, 3]
>>> k_nearest(3, 5.0, s)
[5, 4, 6]