Search code examples
pythonarraysnumpynumpy-ndarrayboolean-indexing

Filter numpy array but keep original value

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?

Solution

  • 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]