Question

假设我有一个numpy数组

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

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

是否有一种简单的方法可以获得以下数组：

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

我尝试过的最明显的方法是[b]，但这并没有产生理想的结果。在此之后，我试着调查masked operations in numpy，但它看起来像是指引我走向错误的方向。

Answer 1

如果a和b都是numpy数组且b严格为1和0：

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

应该注意的是，对于极小的情况，这只会明显加快，并且在范围上比@fattru的回答更加有限：

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

Answer 2

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

使用itertools.compress替代方案：

>>> import itertools
>>> list(itertools.compress(a, b))
[8, 25, 68]

numpy从数组中获取掩码

2 个答案: