Finding closest number index from a vector

Question

If there is a reference vector, for example:

ref = np.array([1., 3., 5.])

And a random vector:

r = np.array([3.1, 4.7, 0.1, 5.5])

Is there a fast way to find the index of the closes number in 'ref' for each number in 'r'?

Expected result is:

[2, 3, 1, 3]

3.1 is closest to 3 so the answer is index 2.

4.7 is closest to 5 so the answer is index 3.

0.1 is closest to 1 so the answer is index 1.

5.5 is closest to 5 so the answer is index 3.

Answer 1

You can do a broadcasting:

np.abs(r[:,None]-ref).argmin(-1)

Output (remember python is 0-indexed):

array([1, 2, 0, 2])