CODEWITHSUNDEEP
pythonnumpy
robintw
robintw

Reputation: 28511

Get indices of array where two conditions (on different arrays) are true

I have two numpy arrays: a and b. I want to select all of the indices where a == 1 and b == 0.

That is, if I have the following arrays:

a = [0, 1, 3, 5, 1, 1, 2]

b = [1, 0, 2, 5, 3, 0, 6]

I would like to get the following indices back:

[1, 5]

How should I do this in numpy? I have tried using the following (suggested by a quick reference guide showing differences between numpy, matlab and IDL):

(a == 1 and b == 0).nonzero()

But that gives an error about truth values being ambiguous.

Any ideas?

Upvotes: 3

Views: 18418

Answers (4)

eat
eat

Reputation: 7530

In []: from numpy import array, logical_and
In []: a= array([0, 1, 3, 5, 1, 1, 2])
In []: b= array([1, 0, 2, 5, 3, 0, 6])
In []: logical_and(a== 1, b== 0).nonzero()[0]
Out[]: array([1, 5])

Obviously this will work as well:

In []: ((a== 1)& (b== 0)).nonzero()[0]
Out[]: array([1, 5])

Upvotes: 14

inspectorG4dget
inspectorG4dget

Reputation: 113905

I'm not an expert with numpy errors, but this is what I would do with regular python lists, and perhaps you could use it here, too:

>>> a = [0, 1, 3, 5, 1, 1, 2]
>>> b = [1, 0, 2, 5, 3, 0, 6]
>>> zip(a, b)
[(0, 1), (1, 0), (3, 2), (5, 5), (1, 3), (1, 0), (2, 6)]
>>> for i, tup in enumerate(zip(a, b)):
...     if tup[0]==1 and tup[1]==0:
...         print i
...
1
5
>>>

Hope this helps

Upvotes: 0

Artsiom Rudzenka
Artsiom Rudzenka

Reputation: 29093

Try to use the following code:

import numpy

a = numpy.array([0, 1, 3, 5, 1, 1, 2])
b = numpy.array([1, 0, 2, 5, 3, 0, 6])

res =  [i for i,v in enumerate(zip(a,b)) if v == (1,0)]

print res

Upvotes: 2

ars
ars

Reputation: 123478

Here's one way:

In [75]: import numpy as np

In [76]: a = np.array([0, 1, 3, 5, 1, 1, 2])

In [77]: b = np.array([1, 0, 2, 5, 3, 0, 6])

In [78]: np.argwhere((a==1) & (b==0)).flatten()
Out[78]: array([1, 5])

Upvotes: 6

Related Questions