How can I apply a 2D mask in numpy but keep the structure of the rows?-CodePudding

I have two numpy arrays e.g.:

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?

CodePudding user response：

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']]