CODEWITHSUNDEEP
pythonnumpy
Salvador Dali
Salvador Dali

Reputation: 222461

numpy get mask from the array

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.

Upvotes: 2

Views: 1457

Answers (2)

Daniel
Daniel

Reputation: 19537

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

Upvotes: 3

falsetru
falsetru

Reputation: 368914

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

Upvotes: 1

Related Questions