Numpy find 1d array elements in 2d array rows

Question

I want to find elements of a 1d array in rows of a 2d array.

Example

In [1]: import numpy as np

In [2]: a = np.array([7,7,7])

In [3]: a
Out[3]: array([7, 7, 7])

In [4]: b = np.arange(15).reshape(3,5)

In [5]: b
Out[5]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

Rows 0 and 2 of b do not have a match for the corresponding element of a, but row 1 has a match in position 2. Expected output:

array([nan, 2, nan])

If there are multiple matches in a given row, the position of the first match should be used.

I can grind out a solution using normal python loops but I'm interested in a way to vectorize this.

What if there are multiple matches in a row? What do you want the result to be? — akuiper
– akuiper, Commented May 25, 2017 at 1:44
What's the expected output if row 1 in b has multiple occurrences of 7? — Spidey
– Spidey, Commented May 25, 2017 at 1:44

Allen Qin · Accepted Answer · 2017-05-25 02:04:11Z

2

Compare b and a element wise, and then find the first True value index for each row, or set to nan if all False.

np.where(np.sum(b==a[:,None],1) > 0, np.argmax(b==a[:,None],1), np.nan)
Out[22]: array([ nan,   2.,  nan])

answered May 25, 2017 at 2:04

Allen Qin

20k9 gold badges55 silver badges68 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Numpy find 1d array elements in 2d array rows

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related