Numpy: give back an array with the n entries closest to a given number

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?

MSeifert · Accepted Answer · 2017-06-13T01:31:36.370

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])

score 0 · Answer 2 · answered Jun 12 '17 at 23:47

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.

Numpy: give back an array with the n entries closest to a given number

2 Answers2