Question

我有一个numpy数组arr，我想创建一个新数组，其中包含旧数组n元素，它们与给定数字x最接近。我在这里找到了一个有用的答案（I have need the N minimum (index) values in a numpy array），但我的代码看起来非常笨拙（我是Python的初学者）：

def give_array_with_closest(x, n, arr):
    newar = np.absolute(arr - (np.ones(len(arr)) * x)) #Subtract x from all array entries and take absolute value, so that the lowest entries are the ones closest to x
    indexar = (newar).argsort()[:n] #get array with indices from n lowest entries of newar
    result = np.empty(n)
    for i in range(n):
        result[i] = arr[indexar[i]]
    return result

由于我对指数不感兴趣而只对实际条目感兴趣，因此在这种情况下，另一个问题的解决方案可能并不是最好的。有没有更有效和更简单的方法来做到这一点？

Answer 1

已经提到你不需要for循环来获取索引的值，你可以简单地使用argsort的结果来索引数组（至少如果您的阵列是1D）。
但您也不需要sort完整数组。您只需使用argpartition即可。可能比排序更快。
另外一点：您可以使用arr - 1之类的矢量化操作。这将从每个元素中减去一个，而无需手动创建新数组（如np.ones(len(arr))）。

将这些放在一起：

def give_array_with_closest(x, n, arr):
    indexarr = np.argpartition(abs(arr - x), n)[:n]
    return arr[indexarr]

测试：

>>> give_array_with_closest(2, 3, np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
array([2, 3, 1])

Answer 2

您可以通过提供索引来创建新的数组一次调用：

def give_array_with_closest2(x, n, arr):
    newar = np.absolute(arr - (np.ones(len(arr)) * x)) 
    return arr[(newar).argsort()[:n]]

所以实际上这是相同的代码，但更优雅，可能更快。

Numpy：返回一个数组，其中n个条目最接近给定的数字

2 个答案: