比较numpy数组中的多个值

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?

Answer 1

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:

• http://scipy-lectures.github.io/intro/numpy/numpy.html#broadcasting
• http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html

Answer 2

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)

Answer 3

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.

Answer 4

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.