Finding nearest element of an array in a particular column of another array

Viewed 148

I want to map one numpy array to another one. My frist array has two columns and thousands of rows:

arr_1 = [[20,  0.5],
         [30, 0.75],
         [40,  1.0],
         [50, 1.25],
         [60,  1.5],
         [70, 1.75],
         ...]

The second array can have a different number of rows and columns:

arr_2 = [[1, 0.45],
         [2, 0.57],
         [4, 0.58],
         [1, 1.69],
         [1, 1.51],
         [1, 0.95],
         ...]

I want to compare the values of the second column of arr_2 with the second column of arr_1 to know which row of arr_2 is closer to which row of arr_1. Then I want to copy the first column of arr_1 into arr_2 from the row with the nearest second column.

For example, 0.45 in arr_2 is closest to 0.5, i.e. first row in arr_1. After finding that, I want to copy the first column of that row (which is 20) into arr_2. The final result would look something like:

arr_2_final = [[1, 0.45, 20],
               [2, 0.57, 20],
               [4, 0.58, 20],
               [1, 1.69, 70],
               [1, 1.51, 60],
               [1, 0.95, 40],
               ...]
2 Answers

Looking up lots of items in an array is easiest done when it is sorted. You can delegate most of the work to np.searchsorted. Since we want to find elements in arr_1, it is the only array that needs to be sorted. I suspect that having a sorted arr_2 will speed things up by reducing the size of the search space for every successive element.

First, find the insertion points where arr_2 would end up in arr_1:

indices = np.searchsorted(arr_1[:, 1], arr_2[:, 1])

Now all you have to do is check for cases where the prior element is closer than the current one. There are two corner cases: when index is 0, you have to accept it, and when it is arr_1.size, you have to take the prior.

indices[indices == arr_1.shape[0]] = arr_1.shape[0] - 1
indices[(indices != 0) & (arr_1[indices, 1] - arr_2[:, 1] > arr_2[:, 1] - arr_1[indices - 1, 1])] -= 1

Doing it in this order saves you the trouble of messing with temporary arrays. The first line ensures that the index arr_1[indices, 1] is always valid. Since index -1 is valid, the second line succeeds as well.

The final result is then

np.concatenate((arr_2, arr_1[indices, 0:1]), axis=1)

If arr_1 is not already sorted, you can do the following:

arr_1 = arr1[np.argsort(arr_1[:, 1]), :]

A quick benchmark shows that on my very moderately powered machine, this approach takes ~300ms for arr_1.shape = (500000, 2) and arr_2.shape = (300000, 2).

I would probably do it this way:

import numpy as np

arr_1= [[20, 0.5], [30, 0.75], [40, 1], [50, 1.25], [60, 1.5], [70, 1.75]]
arr_2= [[1, 0.45], [2, 0.57], [4, 0.58], [1, 1.69], [1, 1.51], [1, 0.95]]

arr_2_np = np.array(arr_2)[:,1]

for row in arr_1:
  idx = np.argmin(np.abs(arr_2_np - row[1]))
  arr_2[idx].append(row[0])

print(arr_2)
Related