CODEWITHSUNDEEP
pythonnumpy
luffe
luffe

Reputation: 1665

Broadcast comparison between two different sized vectors

So I have a vector 1x4 cumpr:

In [119]:

cumpr
Out[119]:
array([ 0.30769231,  0.38461538,  0.53846154,  1.        ])

and a 1x10 vector r:

In [120]:

r
Out[120]:
array([ 0.90737839,  0.05419125,  0.14056591,  0.73179618,  0.34456476,
        0.27682039,  0.74046728,  0.93973422,  0.67263953,  0.44141488])

For every element in r I want to do the following (for example, for the first element in r):

In [86]:

np.nonzero(cumpr > r[0])[0][0]
Out[86]:
3

i.e., return the index of the first element in cumpr where that element is bigger than r[0].

But is there a way to do this for all elements in r without looping? If I try broadcasting

cumpr > r

I just get an error.

Thank you

Upvotes: 0

Views: 71

Answers (1)

DSM
DSM

Reputation: 353419

IIUC, you can simply use np.searchsorted:

>>> np.searchsorted(cumpr, r)
array([3, 0, 0, 3, 1, 0, 3, 3, 3, 2])

Sanity check:

>>> ix = np.searchsorted(cumpr, r)
>>> cumpr[ix]
array([ 1.        ,  0.30769231,  0.30769231,  1.        ,  0.38461538,
        0.30769231,  1.        ,  1.        ,  1.        ,  0.53846154])
>>> cumpr[ix] >= r
array([ True,  True,  True,  True,  True,  True,  True,  True,  True,  True], dtype=bool)

Note though that you don't have "a vector 1x4" and "a 1x10 vector"; you have two one-dimensional arrays, of shape (4,) and (10,). This is important for understanding the broadcasting version:

>>> (cumpr > r[:,None]).argmax(1)
array([3, 0, 0, 3, 1, 0, 3, 3, 3, 2])

which builds an intermediate array by adding a single-element axis.

Upvotes: 3

Related Questions