Suppose I have a numpy array
a = np.array([0, 8, 25, 78, 68, 98, 1])
and a mask array b = [0, 1, 1, 0, 1]
Is there an easy way to get the following array:
[8, 25, 68] - which is first, second and forth element from the original array. Which sounds like a mask for me.
The most obvious way I have tried is a[b], but this does not yield a desirable result. After this I tried to look into masked operations in numpy but it looks like it guides me in the wrong direction.
If a and b are both numpy arrays and b is strictly 1's and 0's:
>>> a[b.astype(np.bool)]
array([ 8, 25, 68])
It should be noted that this is only noticeably faster for extremely small cases, and is much more limited in scope then @falsetru's answer:
a = np.random.randint(0,2,5)
%timeit a[a==1]
100000 loops, best of 3: 4.39 µs per loop
%timeit a[a.astype(np.bool)]
100000 loops, best of 3: 2.44 µs per loop
For the larger case:
a = np.random.randint(0,2,5E6)
%timeit a[a==1]
10 loops, best of 3: 59.6 ms per loop
%timeit a[a.astype(np.bool)]
10 loops, best of 3: 56 ms per loop
>>> a = np.array([0, 8, 25, 78, 68, 98, 1])
>>> b = np.array([0, 1, 1, 0, 1])
>>> a[b == 1]
array([ 8, 25, 68])
Alternative using itertools.compress:
>>> import itertools
>>> list(itertools.compress(a, b))
[8, 25, 68]
