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

23,169
9
62
102

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

Nico Schlömer · Accepted Answer · 2018-02-08T08:51:45.633

searchsorted actually does the trick:

numpy.searchsorted(a, alpha)

The axis argument to argmax helps out; this

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

does the trick. Indeed

import numpy

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

alpha = 1.12

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

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

numpy first occurence in array for an array of reference values

2 Answers2