Python filter numpy array based on mask array

Question

Let's say I have an data numpy array of length N, and a bit mask array of length N.

data = [1,2,3,4,5,6,7,8,9,0]
mask = [0,1,0,1,0,1,0,1,0,1]

Is there a loopless numpy way to create a new array based off data, such that it takes all the entries of data if and only if masks[i] != 0? Like so:

func(data, mask) = [2,4,6,8,0]

Or equivalently in loop notation:

ans = []
for idx in range(mask):
    if mask[idx]:
        ans.append(data[idx])
ans = numpy.array(ans)

Thanks!

Answer 1

You can filter numpy arrays with an array of boolean values. You are starting with an array of integers, which you can't use directly, but you can of course interpret the ones and zeros as booleans and then use it directly as a mask:

import numpy as np

data = np.array([1,2,3,4,5,6,7,8,9,0])
mask = np.array([0,1,0,1,0,1,0,1,0,1])

data[mask.astype(bool)]
# array([2, 4, 6, 8, 0])