ansaurus

Question

Python: Select subset from list based on index set

Answer 1

+3 A:

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

or

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.

KennyTM 2010-07-05 11:32:36

Does using `zip` here introduce a performance penalty?

fuenfundachtzig 2010-07-05 11:36:34

@fuen: Yes. Causes a lot on Python 2 (use [itertools.izip](http://docs.python.org/library/itertools.html#itertools.izip) instead), not so much on Python 3. This is because the `zip` in Python 2 will create a new list, but on Python 3 it will just return a (lazy) generator.

KennyTM 2010-07-05 11:37:56

OK, so I should stick to your 2nd proposal then, because this makes up the central part of my code.

fuenfundachtzig 2010-07-05 11:39:06

@85: why are you worrying about performance? Write what you have to do, if it is slow, then test to find bottlenecks.

PreludeAndFugue 2010-07-05 11:39:22

@PreludeAndFugue: If there are two equivalent options it's good to know which one is faster, and use that one right away.

fuenfundachtzig 2010-07-05 11:42:14

I suspect the second is *slower*, because where did that good_indices list come from in the first place? Probably by enumerating over all of good_objects and saving the indices where good_objects[i] is True. So no savings after all, plus you had to build a second list. Use the first option, with izip in Py2 or zip in Py3, read both lists once, and directly create the desired output with no intermediate lists.

Paul McGuire 2010-07-05 20:29:29

You can just use `from itertools import izip` and use that instead of `zip` in the first example. That creates an iterator, same as Python 3.

Chris B. 2010-07-05 20:34:14

@Paul McGuire: You're right, I'm looping over the properties and applying some tests to figure out which objects are good. So in principle it would be possible to build the lists directly in that loop. This is also probably the fastest way.

fuenfundachtzig 2010-07-05 21:20:36

Answer 2

+2 A:

I see 2 options.

Using numpy:

property_a = numpy.array([545., 656., 5.4, 33.])
property_b = numpy.array([ 1.2,  1.3, 2.3, 0.3])
good_objects = [True, False, False, True]
good_indices = [0, 3]
property_asel = property_a[good_objects]
property_bsel = property_b[good_indices]

Using a list comprehension and zip it:

property_a = [545., 656., 5.4, 33.]
property_b = [ 1.2,  1.3, 2.3, 0.3]
good_objects = [True, False, False, True]
good_indices = [0, 3]
property_asel = [x for x, y in zip(property_a, good_objects) if y]
property_bsel = [property_b[i] for i in good_indices]

WoLpH 2010-07-05 11:34:51

Use 8 spaces to format code within a list.

KennyTM 2010-07-05 11:36:52

Thanks Kenny :)I was wondering why it wasn't working ;)

WoLpH 2010-07-05 11:38:01

Using Numpy is a good suggestion since the OP seems to want to store numbers in lists. A two-dimensional array would be even better.

Philipp 2010-07-05 13:35:56

Answer 3

+2 A:

Use the built in function zip

property_asel = [a for (a, truth) in zip(property_a, good_objects) if truth]

EDIT

Just looking at the new features of 2.7. There is now a function in the itertools module which is similar to the above code.

http://docs.python.org/library/itertools.html#itertools.compress

itertools.compress('ABCDEF', [1,0,1,0,1,1]) =>
  A, C, E, F

PreludeAndFugue 2010-07-05 11:34:56

I'm underwhelmed by the use of `itertools.compress` here. The list comprehension is *far* more readable, without having to dig up what the heck compress is doing.

Paul McGuire 2010-07-05 20:32:06

Hm, I find the code using compress much more readable :) Maybe I'm biased, because it does exactly what I want.

fuenfundachtzig 2010-07-09 15:52:32

ansaurus

tags:

views:

answers:

Python: Select subset from list based on index set

EDIT

related questions