NumPy：布尔ndarray的广播

Question

我有一个问题，我有一个由浮点数组成的ndarray X，假设是100x10的大小。我想对第一列进行某些条件的筛选，并创建一个形状为100x1的布尔类型的ndarray B。然后我想使用B作为索引进入X中提取True所在的值。但是，对于B中的每个True值，我想要提取整行的数据。我认为这应该自动工作，因为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]。为什么广播不能按我期望的方式工作？

- composerMike

---

只需删除更改a形状的行即可。 - Joran Beasley

Answer 1

所谓的“花式索引”的规则已在此处详细说明。特别是，当索引obj是一个bool类型的NumPy数组时，x[obj]

…等同于（但比之更快）x[obj.nonzero()]。 此处的obj.nonzero()返回一个整数索引数组元组（长度为obj.ndim），显示obj中True元素的位置。

因为，

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])

---

如果您想使用布尔数组a选择b的行，那么像Joran Beasley所说的那样，只需将a保持为一维布尔数组即可：

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]]

- unutbu

---

嗯，我的布尔数组b是通过以下方式生成的：b = a[:,1] < 0.5。因此，它将是一个单列的ndarray。有没有办法将其转换为一维数组？ - composerMike

如果 b 的形状为 (n, 1)，那么 np.squeeze(b) 的形状将为 (n,) -- np.squeeze 会从数组的形状中删除单维条目。 - unutbu