логическое маскирование внутренних размеров 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<b>[:, bmask]</b>