How to test whether one of two values present in numpy array or matrix column on single pass?

Question

Take a matrix like the following:

import numpy as np
m = np.matrix([[1,1],
               [2,0],
               [3,1],
               [5,1],
               [5,0]])

Then take two test values:

n1 = 4
n2 = 1

How can I test for both of them (it's guaranteed that only one if any at all will be present) and return that value? Doing two passes is simple enough:

if n1 in m[:, 0]:
    return n1
if n2 in m[:, 0]:
    return n2

What's the best numpy way to consolidate to a single look through m[:, 0]?

There's docs.scipy.org/doc/numpy-1.13.0/reference/generated/…, though I don't know how many passes it actually makes. — user2357112
– user2357112, Commented Oct 15, 2017 at 0:44
I think the best might be something like for value in m[:,0]: if value in [n1, n2]: return value. But I thought there might be something from numpy specifically for this. The numpy.isin shared by @user2357112 doesn't look like it would work as well. — Conner M.
– Conner M., Commented Oct 15, 2017 at 0:48

B. M. · Accepted Answer · 2017-10-15 16:06:38Z

3

If you need simplicity and readability , the simplest can be found with set logic :

{1,4} & set(m[:,0])

Furthermore, the data is actually read exactly one time.

edited Oct 15, 2017 at 16:06

answered Oct 15, 2017 at 14:14

B. M.

18.7k2 gold badges40 silver badges56 bronze badges

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

Comments

John Zwinck · Accepted Answer · 2017-10-15 02:03:53Z

2

You said you want to do it in a single pass through m. You can use isin() (in NumPy 1.13+):

mask = np.isin(m[:,0], [n1,n2])[:,0]

That gives you a boolean mask which is True where the values n1 or n2 are found. To take the first such value:

row = np.where(mask)[0][0]

Of course, that takes a pass through mask which is the same length as m. If you want to optimize this further, you may need to use Numba or Cython to implement a more direct solution using compiled loops.

answered Oct 15, 2017 at 2:03

John Zwinck

252k44 gold badges346 silver badges459 bronze badges

Comments

akilat90 · Accepted Answer · 2017-10-15 02:32:18Z

2

m[:,0][(m[:,0] == n1) | (m[:,0] == n2)][0,0]

Explanation:

m = np.matrix([[1,1],
           [2,0],
           [3,1],
           [5,1],
           [5,0]])
n1 = 4; n2 = 1;

(m[:,0] == n1) returns a boolean matrix for n1's existence

matrix([[False],
    [False],
    [False],
    [False],
    [False]], dtype=bool)

(m[:,0] == n2) returns a boolean matrix for n2's existence

matrix([[ True],
        [False],
        [False],
        [False],
        [False]], dtype=bool)

Since you said that exactly a one of n1 and n2 parameters will be present at a time, |ing the above two will make the indices True for whatever the existing parameter.

(m[:,0] == n1) | (m[:,0] == n2)

matrix([[ True],
        [False],
        [False],
        [False],
        [False]], dtype=bool)

Indexing m[:,0] by the above boolean array,

m[:,0][(m[:,0] == n1) | (m[:,0] == n2)]

matrix([[1]])

We just get the first element out of it

m[:,0][(m[:,0] == n1) | (m[:,0] == n2)][0,0]
1

EDIT:

After numpy 1.13 +, as @John Zwink shows, you can compact the operations up to the last one as np.isin(m[:,0], [n1,n2])[:,0] and then just extract the first element out of it by np.where(np.isin(m[:,0], [n1,n2])[:,0])[0][0]

edited Oct 15, 2017 at 2:32

answered Oct 15, 2017 at 1:59

akilat90

5,7367 gold badges30 silver badges44 bronze badges

1 Comment

kjmerf Over a year ago

The answer works so I don't see why this was voted down +1

Collectives™ on Stack Overflow

How to test whether one of two values present in numpy array or matrix column on single pass?

3 Answers 3

Comments

Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related