python - numpy 从数组中获取掩码

Question

假设我有一个 numpy 数组

a = np.array([0, 8, 25, 78, 68, 98, 1])

和掩码数组 b = [0, 1, 1, 0, 1]

有没有简单的方法得到下面的数组:

[8, 25, 68] - 这是原始数组中的第一个、第二个和第四个元素。这对我来说听起来像是面具。

我尝试过的最明显的方法是 a[b]，但这并没有产生理想的结果。在此之后我试图调查 masked operations in numpy但它似乎将我引向了错误的方向。

Answer 1

如果 a 和 b 都是 numpy 数组并且 b 是严格的 1 和 0:

>>> a[b.astype(np.bool)]
array([ 8, 25, 68])

应该注意的是，这仅在极小的情况下明显更快，并且范围比@falsetru 的回答更受限制:

a = np.random.randint(0,2,5)

%timeit a[a==1]
100000 loops, best of 3: 4.39 µs per loop

%timeit a[a.astype(np.bool)]
100000 loops, best of 3: 2.44 µs per loop

对于较大的情况:

a = np.random.randint(0,2,5E6)

%timeit a[a==1]
10 loops, best of 3: 59.6 ms per loop

%timeit a[a.astype(np.bool)]
10 loops, best of 3: 56 ms per loop