3

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?

robintw
  • 27,571
  • 51
  • 138
  • 205

4 Answers4

14
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])
Tommaso Di Noto
  • 1,208
  • 1
  • 13
  • 24
eat
  • 7,440
  • 1
  • 19
  • 27
6

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])
ars
  • 120,335
  • 23
  • 147
  • 134
2

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
Artsiom Rudzenka
  • 27,895
  • 4
  • 34
  • 52
  • Was what I was in the process of posting :) Nice and concise. – Tom Christie Jul 04 '11 at 11:03
  • Erm, why'd it get downvoted? True, it's not numpy specific, but it is the most pythonic answer. – Tom Christie Jul 04 '11 at 11:32
  • Actually it is always nice to know - what is wrong with code - i mean say me what i am doing wrong - so i will have a chance to not make this error again in the future. – Artsiom Rudzenka Jul 04 '11 at 11:37
  • 1
    @Artsiom Rudzenka: I didn't down vote. FWIW, there is no error per se, and IMHO your solution is very valid when handling pure `python` lists. However `numpy` arrays won't perform very well with this kind of approach. The paradigm is just different. Also, personally, my eyes get hurt when I'm forced to read really needles looping, even it's hidden in list comprehension. IMHO, abstraction levels don't really match with your solution. My 2 cents. Thanks – eat Jul 04 '11 at 16:53
  • I suppose you got downvoted because this solution while correct, doesn't scale to the kind of size that numpy arrays are used for. I didn't test it but I suppose the memory allocations would make your code crash but The numpy-based code should run fine. – LBarret Jul 05 '11 at 07:58
  • @Lionel Barret - thanks. I simply check my code on the given sample - that is why i posted it here. – Artsiom Rudzenka Jul 05 '11 at 08:07
  • I guess it's worth noting that this scales to arbitrarily large arrays fine on python 3, but not on python 2 (under python 3 zip returns an iterator, rather then returning a list). – Tom Christie Jul 05 '11 at 12:25
0

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

inspectorG4dget
  • 110,290
  • 27
  • 149
  • 241