Compare multiple values in numpy array

Question

I have a numpy array

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

I would like to do something like

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

and get

[[False, True,   False, False],
 [False, False,  True,  False],
 [False, False,  False, True ],
 [True,  False,  False, False]]

In other words, I want the ith column to show whether each element of a is equal to i. This feels like the kind of thing that numpy might make easy. Any ideas?

YXD · Accepted Answer · 2013-12-13 23:31:19Z

7

The key concept to use here is broadcasting.

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

The result:

>>> a[..., None] == b[None, ...]
array([[False,  True, False, False],
       [False, False,  True, False],
       [False, False, False,  True],
       [ True, False, False, False]], dtype=bool)

Understanding how to use broadcasting will greatly improve your NumPy code. You can read about it here:

answered Dec 13, 2013 at 23:31

YXD

32.6k15 gold badges79 silver badges117 bronze badges

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

Comments

alko · Accepted Answer · 2013-12-13 23:28:52Z

1

You can reshape to vector and covector and compare:

>>> a = numpy.array([1,2,3,0])
>>> b = numpy.array([0,1,2,3])
>>> a.reshape(-1,1) == b.reshape(1,-1)
array([[False,  True, False, False],
       [False, False,  True, False],
       [False, False, False,  True],
       [ True, False, False, False]], dtype=bool)

answered Dec 13, 2013 at 23:28

alko

48.7k12 gold badges99 silver badges105 bronze badges

Comments

cc7768 · Accepted Answer · 2013-12-13 23:29:32Z

1

The above is one way of doing it. Another possible way (though I'm still not convinced there isn't a better way) is:

import numpy as np
a = np.array([[1, 2, 3, 0]]).T
b = np.array([[0, 1, 2, 3]])
a == b
array([[False,  True, False, False],
   [False, False,  True, False],
   [False, False, False,  True],
   [ True, False, False, False]], dtype=bool)

I think you just need to make sure one is a column vector and one is a row vector and it will do comparison for you.

answered Dec 13, 2013 at 23:29

cc7768

3254 silver badges10 bronze badges

3 Comments

cc7768 Over a year ago

This is the same as the above answer. Beat me haha, I was trying to figure out format. How do I get the arrows to show up?

mechanical_meat Over a year ago

If by "arrows" you mean the REPL's >>> you can copy them from your interpreter.

cc7768 Over a year ago

exactly what I meant thank you. I also learned a new term. Thanks again.

mdml · Accepted Answer · 2013-12-13 23:25:48Z

0

You can use a list comprehension to iterate through each index of a and compare that value to b:

>>> import numpy as np
>>> a = np.array([1,2,3,0])
>>> b = np.array([0,1,2,3])
>>> ans = [ list(a[i] == b) for i in range(len(a)) ]
>>> ans
[[False,  True, False, False],
 [False, False,  True, False],
 [False, False, False,  True],
 [ True, False, False, False]]

I made the output match your example by creating a list of lists, but you could just as easily make your answer a Numpy array.

edited Dec 13, 2013 at 23:25

answered Dec 13, 2013 at 23:19

mdml

23k8 gold badges61 silver badges66 bronze badges

1 Comment

fgregg Over a year ago

Definitely works. I was wondering if there was something more built in to numpy. This feels like the kind of thing that might be part of the library.

Collectives™ on Stack Overflow

Compare multiple values in numpy array

4 Answers 4

Comments

Comments

3 Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

Comments

Comments

3 Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related