python - numpy:广播 bool 值的 ndarray

Question

我有一个情况，我有一个 ndarray X 的 float ，比方说 100x10，我想查看第一列的一些条件并创建一个形状为 100x1 的 bool ndarray B。然后我想使用 B 作为 X 的索引来提取 True 所在的值。然而，对于 B 中的每个 True，我想拉出整行 X。我认为这会自动工作，因为 B 将被广播到 100x10 形状。但是，它似乎无法以这种方式工作。下面是一个使用 2x2 和 2x1 ndarray 的示例。

a = np.array([True, False])
a.shape = (2,1)
b = np.array([1, 2, 3, 4])
b.shape = (2,2)
print(a)
print(b)
print(b[a])

这打印

[[True]
 [False]]

[[ 1 2 ]
 [ 3 4 ]]

 [1]

我希望它打印[1 2]。为什么广播没有按我预期的方式工作？

Answer 1

所谓的“幻想索引”的规则是detailed here .特别是，当索引 obj 是 dtype bool 的 NumPy 数组时，x[obj]

... is always equivalent to (but faster than) x[obj.nonzero()] where, as described above, obj.nonzero() returns a tuple (of length obj.ndim) of integer index arrays showing the True elements of obj.

因为，

In [4]: a.nonzero()
Out[4]: (array([0]), array([0]))

b[a] 等同于 b[a.nonzero()]

In [6]: b[(np.array([0]), np.array([0]))]
Out[6]: array([1])
In [7]: b[a]
Out[7]: array([1])

---

如果您想使用 bool 数组 a 选择 b 的 行，那么，正如 Joran Beasley 所说，只需保留 a 作为一维 bool 数组:

import numpy as np

a = np.array([True, False])
b = np.array([1, 2, 3, 4])
b.shape = (2,2)
print(b[a])    
# [[1 2]]