numpy first occurence in array for an array of reference values

Question

Given a threshold alpha and a numpy array a, there are multiple possibilities for finding the first index i such that arr[i] > alpha; see Numpy first occurrence of value greater than existing value:

numpy.searchsorted(a, alpha)+1
numpy.argmax(a > alpha)

In my case, alpha can be either a scalar or an array of arbitrary shape. I'd like to have a function get_lowest that works in both cases:

alpha = 1.12
arr = numpy.array([0.0, 1.1, 1.2, 3.0])
get_lowest(arr, alpha)  # 2

alpha = numpy.array(1.12, -0.5, 2.7])
arr = numpy.array([0.0, 1.1, 1.2, 3.0])
get_lowest(arr, alpha)  # [2, 0, 3]

Any hints?

Answer 1

You can use broadcasting:

In [9]: arr = array([ 0. ,  1.1,  1.2,  3. ])
In [10]: alpha = array([ 1.12, -0.5 ,  2.7 ])
In [11]: np.argmax(arr > np.atleast_2d(alpha).T, axis=1)
Out[11]: array([2, 0, 3])

To collapse multidimensional arrays, you can use np.squeeze, but you might have to do something special if you want a Python float in your first case:

def get_lowest(arr, alpha):
    b = np.argmax(arr > np.atleast_2d(alpha).T, axis=1)
    b = np.squeeze(b)
    if np.size(b) == 1:
        return float(b)
    return b

Answer 2

searchsorted actually does the trick:

np.searchsorted(a, alpha)

The axis argument to argmax helps out; this

np.argmax(numpy.add.outer(alpha, -a) < 0, axis=-1)

does the trick. Indeed

import numpy as np

a = np.array([0.0, 1.1, 1.2, 3.0])

alpha = 1.12
    
np.argmax(np.add.outer(alpha, -a) < 0, axis=-1)  # 0
np.searchsorted(a, alpha)  # 0

alpha = np.array([1.12, -0.5, 2.7])
np.argmax(np.add.outer(alpha, -a) < 0, axis=-1)  # [2 0 3]
np.searchsorted(a, alpha)  # [2 0 3]