Indexing an array with a logical array of the same shape in Python

Question

I would like to be able to do:

a = [1,2,3]
a[[True, True, False]]
>> array([1,2])

I just can't find the simple way how... Thanks!

Answer 1

There's itertools.compress :

>>> from itertools import compress
>>> a = [1,2,3]
>>> mask = [True, True, False]
>>> list(compress(a, mask))
[1, 2]

---

If you're using numpy , you can slice directly with the mask:

>>> np.array(a)[np.array(mask)]
>>> array([1, 2])

Answer 2

If a and the mask are truly Python arrays, then you can do a simple list comprehension:

arr = [1,2,3]
mask = [True, True, False]
[a for a, m in zip(arr, mask) if m]

If you are OK with additional imports, you can use @moses-koledoye's suggestion of using itertools.compress .

If on the other hand you are using numpy , as the final output of array([1,2]) suggests, you can just do the indexing directly:

import numpy as np
arr = np.array([1, 2, 3])
mask = np.array([True, True, False])
arr[mask]

Note that mask has to be an actual np.boolean array. You can not just use a Python list like [True, True, False] . This is because np.array.__getitem__ checks if the input is exactly another array. If not, the input is converted to integers, so you end up effectively indexing with [1, 1, 0] instead of a mask. You can get a lot more details on this particular tangential issue here: https://stackoverflow.com/a/39168021/2988730

Answer 3

You can use zip and list comprehension.

a = [1,2,3]
b = [True, True, False]

c = [i for i,j in zip(a,b) if j]

Answer 4

You can play with built-in filter and map

k = [True, True, False, True]
l = [1,2,3,4]

print map(lambda j: j[1], filter(lambda i: i[0], zip(k,l)))