0

I need to use searchsorted() along the first axis of an n dimensional array against an n-1 dimensional array. This is illustrated for a 3 dimensional array:

def find(x,y):
    rows = np.zeros(y.shape, dtype=np.int)
    for ii in range(y.shape[0]):
        for jj in range(y.shape[1]):
            rows[ii, jj] = np.searchsorted(x[:, ii, jj], y[ii, jj])
    return rows


np.random.seed(42)
h = np.random.random((10,2,10))
h.sort(axis=0)
f = np.random.random((2,10))
print(find(h,f))

result is something like this:

[[10  0  3 10  5  0  1  5  0  6]
 [ 7  9  2  1 10  1  3  9  5 10]]

Obviously using those loops is probably slow and this solution does not extend to n dimensions. I am concerned with speed and searchsorted() is very fast on sorted arrays when compared to things like argmax() and where().

Is there some way to apply searchsorted() along an axis like this without looping? The utility apply_along_axis kind of works, but makes a big array in 2 dimensions that you need to grab the diagonal of, which seems like a waste. I can't tell what it is doing in more than 2 dimensions.

Chebyshev
  • 71
  • 1
  • 5
  • In `find()` you assume that the first axis of `x` is sorted. However, with `h.sort()` the default will sort the array along the last axis. You should use `h.sort(axis=0)`. – Marc Jan 21 '21 at 20:14
  • I think your dimensions may not be right. Did you mean `h.sort(0)`? This would actually be faster if your sorted dimension was the last dimension, not the first. – Mad Physicist Jan 21 '21 at 20:28
  • I found a similar discussion: [vectorized-searchsorted-numpy](https://stackoverflow.com/questions/40588403/vectorized-searchsorted-numpy) – Marc Jan 21 '21 at 20:29
  • @Marc. Excellent find. I was working on a solution that's exactly along the lines of what Divakar suggested – Mad Physicist Jan 21 '21 at 20:30
  • Yes, should have been `sort(axis=0)`, but that's not really the point of the question. The linked discussion vectorizes the 2D solution, but doesn't address a nD solution. – Chebyshev Jan 21 '21 at 22:38

0 Answers0