# 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)?

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.

Categories: questions
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.