如何在NumPy中应用2D掩码并生成保留行结构的列表数组？

问题背景 #

现有两个NumPy数组：

import numpy as np

names = np.array(['A', 'B', 'C', 'D'])        # 形状 (B,)
mask = np.array([                             # 形状 (N, B)
    [True, False, False, False],
    [False, True, False, True ],
    [True, True, False, False]
])

需要生成一个一维结构，其中每个元素是mask对应行中True位置所对应的names元素组成的列表，预期结果（注：原问题中第二行预期结果应为['B', 'D']，与mask逻辑匹配）：

result = [
    ['A'],
    ['B', 'D'],
    ['A', 'B']
]

解决方案 #

方法1：列表推导式（最直观高效）

利用布尔索引结合列表推导式，遍历mask的每一行提取对应元素：

result = [names[row].tolist() for row in mask]

• 逻辑：对mask的每一行row，names[row]会自动筛选出row中True对应的names元素，再通过.tolist()转为Python列表。
• 优势：代码简洁，处理不等长元素场景适配性好，无需额外处理NumPy的数组形状限制。

方法2：生成NumPy object类型数组

如果需要保持NumPy数组格式，可以生成dtype=object的数组（存储Python列表）：

result_arr = np.array([names[row] for row in mask], dtype=object)

输出结果为：

array([list(['A']), list(['B', 'D']), list(['A', 'B'])], dtype=object)

• 注意：这种数组无法直接进行常规的NumPy数值运算，仅适合需要保持数组容器形式的场景。

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