Given a list and a bitmask, how do I return the values at the indices that are True?

Question

I start with the following list s and bitmask b:

s = ['baa', 'baa', 'black', 'sheep', 'have', 'you', 'any', 'wool']
b = [1, 0, 0, 0, 1, 1, 1, 0] # or any iterable with boolean values

How do I write some function apply_bitmask(s, b) so that it returns

['baa', 'have', 'you', 'any']

Answer 1

Python 3.1 itertools.compress (or Python 2.7's if you haven't upgraded yet) does exactly that (the list comprehension is a real close second):

import itertools
filtered = itertools.compress(s, b)

Note that this produces an iterator, not a list. Saves memory, but if you need to iterate it several times or use indices, you can always use list(itertools.compress(s, b)). Still shorter.

Answer 2

[ item for item, flag in zip( s, b ) if flag == 1 ]

Answer 3

You can use list comprehensions:

newList = [word for (word, mask) in zip(s,b) if mask]
# Note: Could also use 'if mask == blah', if mask is not a boolean-compatible type.

This first takes the original two lists, and zips them together, so that you get a (temporary - this is still inside the list comp!) list of pairs of words and their masks - something like [('baa',1), ('baa',0),...]. Then only the words that have a mask of 1 (if mask == 1) are added to the newList.

Answer 4

Another take on list comprehension, witout using zip

newList = [item for i, item in enumerate(s) if b[i]]