How to filter 3D array with a 2D mask-CodePudding

I have a (m,n,3) array data and I want to filter its values with a (m,n) mask to receive a (x,3) output array.

The code below works, but how can I replace the for loop with a more efficient alternative?

import numpy as np

data = np.array([
    [[11, 12, 13], [14, 15, 16], [17, 18, 19]],
    [[21, 22, 13], [24, 25, 26], [27, 28, 29]],
    [[31, 32, 33], [34, 35, 36], [37, 38, 39]],
])
mask = np.array([
    [False, False, True],
    [False, True, False],
    [True, True, False],
])

output = []
for i in range(len(mask)):
    for j in range(len(mask[i])):
        if mask[i][j] == True:
            output.append(data[i][j])
output = np.array(output)

The expected output is

np.array([[17, 18, 19], [24, 25, 26], [31, 32, 33], [34, 35, 36]])

CodePudding user response：

import numpy as np

data = np.array([
    [[11, 12, 13], [14, 15, 16], [17, 18, 19]],
    [[21, 22, 13], [24, 25, 26], [27, 28, 29]],
    [[31, 32, 33], [34, 35, 36], [37, 38, 39]],
])
mask = np.array([
    [False, False, True],
    [False, True, False],
    [True, True, False],
])

output = data[mask]