3d矩阵中的2d布尔选择

Question

In a related question我了解到如果我有一个形状MxMxN的数组，并且我想根据形状MxM的布尔矩阵进行选择，我可以简单地做

data[select, ...]

并完成它。不幸的是，现在我的数据顺序不同了：

import numpy as np
data = np.arange(36).reshape((3, 4, 3))
select = np.random.choice([0, 1], size=9).reshape((3, 3)).astype(bool)

对于data索引i0, i1, i2中的每个元素，如果select[i0, i2] == True，则应选择该元素。

如何继续我的选择，而不必像

那样做一些低效率的事情

data.flatten()[np.repeat(select[:, None, :], 4, axis=1).flatten()]

Answer 1

一种方法是简单地使用np.broadcast_to进行广播而不进行实际复制，并直接使用该广播掩码来屏蔽所需的元素 -

mask = np.broadcast_to(select[:,None,:], data.shape)
out = data[mask]

另一种方式，可能更快的方法是获取索引，然后用那些索引。由此获得的元素将按axis=1排序。实现看起来像这样 -

idx = np.argwhere(select)
out = data[idx[:,0], :, idx[:,1]]