如何在numpy中将2D mask应用于1D数组以输出n×2格式的结果数组

numpy 没有完全匹配该需求的单一内置方法，但可以通过几个基础内置函数组合实现，步骤非常简洁。

首先我们需要用np.where提取mask中所有True值的行、列坐标，再用两组坐标分别索引1D数组，最后按列拼接即可得到目标结果。

完整实现代码如下：

import numpy as np

numbers = np.array([ 111, 2298, 5820, 6078, 7250])
mask = np.array([[False, False, False, False, False],
                 [ True, False, False, False, False],
                 [False, False, False, False, False],
                 [False, False, False, False, False],
                 [ True,  True, False, False, False]])

# 提取mask中True的行、列索引
row_idx, col_idx = np.where(mask)
# 用索引取数后拼接为n×2数组
res = np.column_stack([numbers[row_idx], numbers[col_idx]])

运行后输出结果为：

array([[2298,  111],
       [7250,  111],
       [7250, 2298]])

如果需要调整每组两个值的顺序，仅需要调换column_stack中两个数组的位置即可。

用到的两个核心函数均为numpy内置方法：

• np.where()：传入布尔数组时，返回所有值为True的位置的各维度索引数组
• np.column_stack()：将多个等长1D数组按列拼接为2D数组

内容的提问来源于stack exchange，提问作者Lukas S