假设我有一个 numpy 数组
a = np.array([0, 8, 25, 78, 68, 98, 1])
和一个掩码数组b = [0, 1, 1, 0, 1]
有没有一种简单的方法来获取以下数组:
[8, 25, 68]
- 这是原始数组中的第一个、第二个和第四个元素。这对我来说听起来像是一个面具。
我尝试过的最明显的方法是 a[b],但这并没有产生理想的结果。在此之后,我尝试查看numpy 中的掩码操作,但它看起来像是将我引导到错误的方向。
假设我有一个 numpy 数组
a = np.array([0, 8, 25, 78, 68, 98, 1])
和一个掩码数组b = [0, 1, 1, 0, 1]
有没有一种简单的方法来获取以下数组:
[8, 25, 68]
- 这是原始数组中的第一个、第二个和第四个元素。这对我来说听起来像是一个面具。
我尝试过的最明显的方法是 a[b],但这并没有产生理想的结果。在此之后,我尝试查看numpy 中的掩码操作,但它看起来像是将我引导到错误的方向。
如果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
>>> 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]