使用相同掩码的 ndarray 的布尔掩码内部维度

Question

I have a 3d ndarray and I want to remove all values of lower dimensions given the same mask efficiently and cleanly. Hopefully code below explains what I'm looking for better:

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