基于掩码数组的Python过滤numpy数组

Question

假设我有一个长度为 N 的数据 numpy 数组和一个长度为 N 的位掩码数组。

data = [1,2,3,4,5,6,7,8,9,0]
mask = [0,1,0,1,0,1,0,1,0,1]

是否有一种无环的 numpy 方法可以基于数据创建一个新数组，这样当且仅当 mask[i] != 0 时它才会获取所有数据条目？像这样：

func(data, mask) = [2,4,6,8,0]

或者等价于循环符号：

ans = []
for idx in range(mask):
    if mask[idx]:
        ans.append(data[idx])
ans = numpy.array(ans)

谢谢！

Answer 1

您可以使用布尔值数组过滤 numpy 数组。您从一个不能直接使用的整数数组开始，但您当然可以将 1 和 0 解释为布尔值，然后然后将其直接用作掩码：

import numpy as np

data = np.array([1,2,3,4,5,6,7,8,9,0])
mask = np.array([0,1,0,1,0,1,0,1,0,1])

data[mask.astype(bool)]
# array([2, 4, 6, 8, 0])