Filter numpy array but keep original value

Question

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?

Answer 1

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]

Answer 2

Another idea would be to initialize the result array c with zeroes (based on the shape of a ) and then use a boolean mask to fill in the values from a . Below is an illustration:

# result array
In [28]: c = np.zeros_like(a)  

# get the indices where nonzero values exist in array `b`
In [29]: nonzero_idx = np.nonzero(b)  

# extract the corresponding values from the array `a` and
# fill in the array `c` using the same indices
In [30]: c[nonzero_idx] = a[nonzero_idx] 

In [31]: c 
Out[31]: array([1, 0, 0, 0, 0, 4])

---

The explanation for the case of using numpy.where :

In [42]: np.where(b >  0, *[a, 0])   
Out[42]: array([1, 0, 0, 0, 0, 4])

the part b > 0 is the condition that we're gonna check, with the entries in the array a . If this condition is satisfied, then those elements would be returned, else we a 0 would be returned; this is why we pass 0 as a third argument to numpy.where()

If you want some other values to be filled in instead of 0 s, then pass the desired value. For example, let's say we want a -9 to be filled in the places where the condition ( b > 0 ) is not satisfied (ie False ), then we would write something like:

In [43]: np.where(b >  0, *[a, -9])     
Out[43]: array([ 1, -9, -9, -9, -9,  4])