python - 过滤 numpy 数组但保留原始值

Question

换句话说，numpy 是否支持“筛选”？

我有两个数组:

a = np.array([1, 0, 2, 3, 0, 4])
b = np.array([1, 0, 0, 0, 0, 6])

我想要的是返回一个新数组 c，它包含基于 b 掩码的 a 的原始值:

c = a[b > 0] 
>> c
np.array([1, 4])
# but what I want is:
# np.array([1, 0, 0, 0, 0, 4])

我可以通过列表理解来解决这个问题:

c = [a[i] if b[i] > 0 else 0 for i in range(len(a))]

我也可以制作一个 mask ，但这需要 2 次迭代:

mask = [1 if b[i] > 0 else 0 for i in range(len(b))]  
c = ma.array(a, mask=mask)

numpy 是否有一些原生的东西允许一个数组对另一个数组充当筛子？

Answer 1

使用np.where :

result = np.where(b > 0, a, 0)
print(result)

或者只是相乘:

import numpy as np

a = np.array([1, 0, 2, 3, 0, 4])
b = np.array([1, 0, 0, 0, 0, 6])

print(a * (b > 0))

输出

[1 0 0 0 0 4]