In other words, does numpy support 'sieving'?
I have two arrays:
a = np.array([1, 0, 2, 3, 0, 4])
b = np.array([1, 0, 0, 0, 0, 6])
What I want is to return a new array, c
, that contains the original values of a
based on the mask of b
:
c = a[b > 0]
>> c
np.array([1, 4])
# but what I want is:
# np.array([1, 0, 0, 0, 0, 4])
I can get around this by doing a list comprehension:
c = [a[i] if b[i] > 0 else 0 for i in range(len(a))]
I could also make a mask, but that would require 2x iteration:
mask = [1 if b[i] > 0 else 0 for i in range(len(b))]
c = ma.array(a, mask=mask)
Is there something native to numpy that allows one array to act like a sieve for another array?
Use np.where:
result = np.where(b > 0, a, 0)
print(result)
Or just multiply:
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))
Output
[1 0 0 0 0 4]