Python:根据索引集选择列表子集

136 投票
5 回答
300814 浏览
提问于 2025-04-16 00:50

我有几个列表,它们的长度都是一样的(每个列表都指定了一个对象的属性):

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_aselproperty_bsel 等等,里面只包含那些被 True 标记的值或者索引的值,最简单的方法是什么呢?

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

5 个回答

18

使用内置的函数 zip

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

编辑

刚刚看了2.7的新特性。现在在itertools模块中有一个函数,和上面的代码很相似。

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

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

我看到有两个选择。

  1. 使用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]
    
  2. 使用列表推导式并结合zip函数:

    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]
    
186

你可以直接使用列表推导式

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]

后者的速度更快,因为good_indices的数量比property_a的长度要少,前提是good_indices是提前计算好的,而不是实时生成的。


编辑:第一个选项相当于从Python 2.7/3.1开始提供的itertools.compress。可以查看@Gary Kerr的回答。

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

撰写回答