如何在 numpy 中应用 2D 蒙版但保持行的结构？

Question

I have two numpy arrays eg:

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

I want to get an 1 dimensional array of lists with shape (N,) , where for each row in mask I have a list of the names in names for which the value was True . Here that would be:

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

Any idea how?

Answer 1

For each array m in mask , you can subset names based on the values of m .

[list(names[m]) for m in mask]

# [['A'], ['B', 'D'], ['A', 'B']]