获取任意形状的多维numpy数组的索引，并将其用于另一个数组

如何解决获取任意形状的多维numpy数组的索引，并将其用于另一个数组

我有两个形状相同的多维数组a和b：

print(a.shape) -> (100,20,3,3)
print(b.shape) -> (100,3)

现在，我只想从a中获取符合特定条件的值，例如：

wanted_value = 5
result_a = a[a[:,:,0] == wanted_value]

这很好，但是，现在我想从另一个索引数组b中获取数据，该数组与数组a中找到的元素的索引相对应。是否可以以某种方式使用从搜索数组a到数组b上的某个条件所获得的索引，以从数组b中获取与相同形状的相同索引对应的数据该数组a有？我之所以需要这样做，是因为我无法搜索数组b的条件，但数组b的数据/元素分别对应于数组a。

这可能吗？我尝试过这样的事情：

indices = np.where(a[a[:,0] == wanted_value])
results_b = b[indices]

解决方法

是的。它比您预期的还要简单：

b[a[:,:,0] == wanted_value]

它的工作方式是a[:,0] == wanted_value本身会使用dtype=bool创建一个Numpy数组，在文档中称为 mask 。您可以将其应用于任何形状相同的数组。

Numpy支持布尔索引，即当您使用布尔数组对numpy数组进行索引时，将返回所有真值为True的值。当obj是布尔类型的数组对象时，将发生这种高级索引。比较运算符返回布尔对象，因此可以将它们用作索引。

a = np.random.randint(0,500,(100,20,3,3))
b = np.copy(a)
assert np.array_equal(a[a[:,0] == 1],b[a[:,0] == 1])

从here检查官方文档。