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

Question

I have a numpy array arr and I want to create a new array with the n elements of the old one, that are the closest to a given number x . I've found a helpful answer here ( I have need the N minimum (index) values in a numpy array ), but my code seems very clumsy (I'm a beginner in 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

Since I'm not interested in the indices but only the actual entries, maybe the solution of the other question isn't the best in this case. Is there a more efficient and simpler way to do this?

Answer 1

• It has been mentioned already that you don't need the for-loop to get the values for the indices, you can simply use the result of argsort to index the array (at least if your arrays are 1D).

• But you also don't need to sort the full array. You could simply use argpartition . That could be faster than sorting.

• As additional point: You can use vectorized operations like arr - 1 . That will subtract one from each element without needing to create a new array manually (like your np.ones(len(arr)) ).

So putting these together:

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

And the test:

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

Answer 2

You could create new array one-call by providing indices:

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

So actually this is the same code, but more elegant and probably faster.