python - 如何使用 2D 蒙版过滤 3D 阵列

Question

我有一个 (m,n,3) 数组 data，我想用 (m,n) 过滤它的值用于接收 (x,3) output 数组的掩码。

下面的代码可以工作，但我怎样才能用更有效的替代方法替换 for 循环？

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)

预期的输出是

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

Answer 1

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]