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)

对于data索引中的每个元素i0, i1, i2,应该选择if 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]]