Avoiding for-loop in NumPy 1D nearest neighbors

Question

I have the following code in which I get the N nearest neighbors in 1D:

import numpy as np

def find_nnearest(arr, val, N):
    idxs = []
    for v in val:
        idx = np.abs(arr - v).argsort()[:N] 
        idxs.append(idx)
    return np.array(idxs)

A = np.arange(10, 20)
test = find_nnearest(A, A, 3)
print(test)

which clearly uses a for-loop to grab the idx's. Is there a numpythonic way to avoid this for-loop (but return the same array)?

Answer 1

To answer your question, yes, you can make the code more "numpythonic" by avoiding the explicit for-loop. You can use broadcasting to compute the absolute differences between the entire array arr and the array val in a single operation. Here's an updated version of your code:

import numpy as np

def find_nnearest(arr, val, N):
    idxs = np.abs(arr[:, None] - val).argsort(axis=0)[:N]
    return idxs

A = np.arange(10, 20)
test = find_nnearest(A, A, 3)
print(test)

In this version, arr[:, None] introduces a new axis to arr, making it a 2D array where each row corresponds to an element in arr. The broadcasting then takes care of computing the absolute differences for all combinations of elements between arr and val in a single step. The rest of the code remains similar, but now the indexing is done using axis=0 to specify that the sorting should be done along the new axis.