Filter elements from list based on True/False from another list

user_326389 0
Viewed 1025

Is there an idiomatic way to mask elements of an array in vanilla Python 3? For example:

a = [True, False, True, False]
b = [2, 3, 5, 7]
b[a]

I was hoping b[a] would return [2, 5], but I get an error:

TypeError: list indices must be integers or slices, not list

In R, this works as I expected (using c() instead of [] to create the lists). I know NumPy has MaskedArray that can do this, I'm looking for an idiomatic way to do this in plain vanilla Python. Of course, I could use a loop and iterate through the mask list and the element list, but I'm hoping there's a more efficient way to mask elements using a higher level abstraction.

masked-array python-3.x list python
5 Answers

You can use itertools.compress:

>>> from itertools import compress
>>> a = [True, False, True, False]
>>> b = [2, 3, 5, 7]


>>> list(compress(b, a))
[2, 5]

Refer "itertools.compress()" document for more details

user_20633610

I think there aren't a lot of ways to do this, you could use a zip:

print([y for x, y in zip(a, b) if x])

Output:

[2, 5]

You could also create a class with __getitem__ for this:

class index:
    def __init__(self, seq):
        self.seq = seq
    def __getitem__(self, boolseq):
        return [x for x, y in zip(boolseq, self.seq) if y]
print(index(a)[b])

Output:

[2, 5]
user_87083640

You could use a list-comprehension:

[b[i] for i in range(len(a)) if a[i]]

or

[b[i] for i,mask in enumerate(a) if maks]

In both cases a list is created by iterating over each element and only inserting it if the mask is true.

user_81065830

You can use a pandas.Series which allows, like for dataframe, to filter data with a boolean array

from pandas import Series

a = [True, False, True, False]
b = [2, 3, 5, 7]

res = Series(b)[a].tolist()

print(res)  # [2, 5]
user_72126860

Among the many options you can:

1. Use itertools compress

from itertools import compress
a = [True, False, True, False]
b = [2, 3, 5, 7]

result_itertools = list(compress(b, a))
print(result_itertools)

2. Use Filter Function

result_filter = list(filter(lambda x: x[0], zip(a, b)))
for item in result_filter:
    print(item[1])
# 2
# 5

3. Use List Comprehension

result_comprehension = [value for bool_, value in zip(a, b) if bool_]
print(result_comprehension)
# [2, 5]
user_139039420
Related