使用相同的掩码对ndarray的内部掩码进行布尔掩码

Question

我有一个3d ndarray，我想有效，干净地删除给定相同蒙版的所有较小尺寸的值。希望下面的代码解释我正在寻找的更好的东西：

arr   = np.random.randint(0,3, (5, 3,4))
bmask = arr[0] > 0 

#######
#Naive way
########
result = []
for i in range(arr.shape[0]):
    result.append(arr[i][bmask])
result = np.stack(result)

####### 
# Better Way
#######
better_way = ???????


assert np.all(result == better_way)

Answer 1

您可以在第一个维度上使用冒号，而在其余维度上使用bmask，例如：

better_way = arr[:, bmask]