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)?
>Solution :
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.