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?

Just to clarify: `arr` is always 1d and sorted? And you treat the elements of `alpha` independently? What's the relative size of `arr` and `alpha`? A simple iteration on `alpha` might be the fastest, especially if `searchsorted` is the optimal tool for a single value. — hpaulj, Jan 21 '18 at 18:03
@hpaulj `alpha` can be an array of arbitrary shape, and is typically much larger than `arr`. — Nico Schlömer, Jan 21 '18 at 21:02
I'm reminded of a `padding` problem - with variable length lists. @Divakar's style of solution might be useful, https://stackoverflow.com/questions/43924187/load-data-with-rows-of-different-sizes-into-numpy-array — hpaulj, Jan 21 '18 at 21:08

xnx · Answer 1 · 2018-01-21T14:25:16.573

2

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

edited Jan 21 '18 at 14:25

answered Jan 21 '18 at 14:15

xnx

24,509
11
70
109

Thanks for the answer! Unfortunately, this returns an array (of shape `1,1`) if the input is a float. I need a float. – Nico Schlömer Jan 21 '18 at 14:16
Easily fixed @NicoSchlömer, no? – jtlz2 Nov 14 '22 at 11:05

Nico Schlömer · Accepted Answer · 2022-11-14T15:35:30.803

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]

edited Nov 14 '22 at 15:35

answered Jan 21 '18 at 14:18

Nico Schlömer

53,797
27
201
249

searchsorted didn't work for me - but the argmax / outer route did. Probably because not sorted :\ – jtlz2 Nov 14 '22 at 11:09

numpy first occurence in array for an array of reference values

2 Answers2