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

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)

对于索引为i0, i1, i2 data每个元素，如果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]]