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