Search code examples
pythonarrayslistmasking

Python: Elegant and efficient ways to mask a list


Example:

from __future__ import division
import numpy as np

n = 8
"""masking lists"""
lst = range(n)
print lst

# the mask (filter)
msk = [(el>3) and (el<=6) for el in lst]
print msk

# use of the mask
print [lst[i] for i in xrange(len(lst)) if msk[i]]

"""masking arrays"""
ary = np.arange(n)
print ary

# the mask (filter)
msk = (ary>3)&(ary<=6)
print msk

# use of the mask
print ary[msk]                          # very elegant  

and the results are:

>>> 
[0, 1, 2, 3, 4, 5, 6, 7]
[False, False, False, False, True, True, True, False]
[4, 5, 6]
[0 1 2 3 4 5 6 7]
[False False False False  True  True  True False]
[4 5 6]

As you see the operation of masking on array is more elegant compared to list. If you try to use the array masking scheme on list you'll get an error:

>>> lst[msk]
Traceback (most recent call last):
  File "<interactive input>", line 1, in <module>
TypeError: only integer arrays with one element can be converted to an index

The question is to find an elegant masking for lists.

Updates:
The answer by jamylak was accepted for introducing compress however the points mentioned by Joel Cornett made the solution complete to a desired form of my interest.

>>> mlist = MaskableList
>>> mlist(lst)[msk]
>>> [4, 5, 6]

Solution

  • If you are using numpy:

    >>> import numpy as np
    >>> a = np.arange(8)
    >>> mask = np.array([False, False, False, False, True, True, True, False], dtype=np.bool)
    >>> a[mask]
    array([4, 5, 6])
    

    If you are not using numpy you are looking for itertools.compress

    >>> from itertools import compress
    >>> a = range(8)
    >>> mask = [False, False, False, False, True, True, True, False]
    >>> list(compress(a, mask))
    [4, 5, 6]