Pythonic方法：依据二进制掩码从MxNxD三维数组提取切片

没问题！这是NumPy里非常常见的索引场景，完全不需要写繁琐的循环，用高级索引就能一步实现，而且效率拉满。

核心思路 #

NumPy支持直接用掩码的非零位置索引来提取多维数组的元素：当你用np.nonzero(M)获取到掩码中1的行、列索引后，直接把这两组索引传入原数组I的索引位，就能自动得到对应位置的D维向量，最终组合成k×D的结果数组。

完整代码示例 #

import numpy as np

# 先构造示例数据方便测试
mask = np.array([[1, 0, 1], [0, 1, 0]])  # 2×3的二进制掩码，共有3个1（k=3）
input_array = np.random.rand(2, 3, 4)  # 2×3×4的原数组，D=4

# 获取掩码非零位置的行、列索引
row_indices, col_indices = np.nonzero(mask)

# 提取对应元素，直接得到k×D的结果数组
extracted_array = input_array[row_indices, col_indices]

print(extracted_array.shape)  # 输出 (3, 4)，完全符合需求

更简洁的写法 #

如果你追求极致简洁，甚至可以把np.nonzero直接嵌入索引操作里，一行代码搞定：

extracted_array = input_array[np.nonzero(mask)]

为什么高效？ #

这种方式依赖NumPy底层的C优化实现，比手动写Python循环快几个数量级，尤其是当你的数组规模很大（M、N、D数值较高）时，性能优势会特别明显。

内容的提问来源于stack exchange，提问作者Felix Goldberg