fuenfundachtzig fuenfundachtzig - 3 months ago 12x
Python Question

Python: Select subset from list based on index set

I have several lists having all the same number of entries (each specifying an object property):

property_a = [545., 656., 5.4, 33.]
property_b = [ 1.2, 1.3, 2.3, 0.3]

and list with flags of the same length

good_objects = [True, False, False, True]

(which could easily be substituted with an equivalent index list:

good_indices = [0, 3]

What is the easiest way to generate new lists
, ... which contain only the values indicated either by the
entries or the indices?

property_asel = [545., 33.]
property_bsel = [ 1.2, 0.3]


You could just use list comprehension:

property_asel = [val for is_good, val in zip(good_objects, property_a) if is_good]


property_asel = [property_a[i] for i in good_indices]

The latter one is faster because there are fewer good_indices than the length of property_a, assuming good_indices are precomputed instead of generated on-the-fly.

Edit: The first option is equivalent to itertools.compress available since Python 2.7/3.1. See @Gary Kerr's answer.

property_asel = list(itertools.compress(good_objects, property_a))