CODEWITHSUNDEEP
pythonarraysnumpy
fgregg
fgregg

Reputation: 3249

Compare multiple values in numpy array

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?

Upvotes: 4

Views: 4493

Answers (4)

YXD
YXD

Reputation: 32511

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:

Upvotes: 7

cc7768
cc7768

Reputation: 325

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.

Upvotes: 1

alko
alko

Reputation: 48307

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)

Upvotes: 1

mdml
mdml

Reputation: 22882

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.

Upvotes: 0

Related Questions