python numpy.where function for lists

Question

I have a python program with two lists that look like the example below:

list1=["a","b","c","d"]
list2=[0,1,1,0]

Is there a elegant way to create a third list, that contains the elements of list1 at the position where list2 is 1? I am looking for something similar to the numpy.where function for arrays or better the elegant way:

array1=numpy.array(["a","b","c","d"])
array2=numpy.array([0,1,1,0])
array3=array1[array2==1]

Is it possible to create a list3 equivalent to array3, containing in this example "b" and "c" or do I have to cast or to use a loop?

Answer 1

You could use a list comprehension here.

>>> array1 = ["a", "b", "c", "d"]
>>> array2 = [0, 1, 1, 0]
>>> [array1[index] for index, val in enumerate(array2) if val == 1] # Or if val
['b', 'c']

Or use,

>>> [a for a, b in zip(array1, array2) if b]
['b', 'c']

Answer 2

This is exactly what itertools.compress does.

>>> list1=["a","b","c","d"]
>>> list2=[0,1,1,0]
>>> import itertools
>>> list(itertools.compress(list1, list2))
['b', 'c']