python - 比较 numpy 数组中的多个值

Question

我有一个 numpy 数组

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

我想做类似的事情

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

得到

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

换句话说，我希望第i列显示a的每个元素是否等于i。这感觉就像是 numpy 可能会变得容易的事情。有什么想法吗？

Answer 1

这里使用的关键概念是广播。

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

结果:

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

了解如何使用广播将大大改进您的 NumPy 代码。你可以在这里阅读:

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