如何在Python中有效地提取列表元素的特定子集

Question

我有一个名为vals的float的长度N列表和名为bits的0和1的长度N列表。我想提取两个列表：vals中与bits中的0相对应的元素，以及与bits中的1相对应的其余元素。我目前正在做：

valsbits = zip(vals,bits)

els0 = [v for v,b in valsbits if b == 0]

els1 = [v for v,b in valsbits if b == 1]

但一定有更好的办法。另外，我对许多不同的bits向量执行此操作，因此可能有一个聪明的方法来完成整个操作。在

Answer 1

您可以使用^{}，它在选择器中生成与true相对应的元素。在

但是，这将需要复制bits并反转一个副本以选择零的元素，最终结果是：

from operator import not_
true_values = list(compress(sequence, bits))
false_values = list(compress(sequence, map(not_, bits)))

---

我相信使用一个简单的for循环会更容易和更快，因为它只进行一次迭代：

^{pr2}$

---

出于好奇，以下是一些包含各种解决方案的微基准测试：

In [12]: import random

In [13]: value = 'a' * 17000

In [14]: selectors = [random.randint(0, 1) for _ in range(17000)]

In [15]: %%timeit
    ...: true_values = [v for v,b in zip(value, selectors) if b == 1]
    ...: false_values = [v for v,b in zip(value, selectors) if b == 0]
    ...: 
100 loops, best of 3: 2.56 ms per loop

In [16]: %%timeit
    ...: true_values = []
    ...: false_values = []
    ...: for bit,val in zip(selectors, value):
    ...:     if bit:
    ...:         true_values.append(val)
    ...:     else:
    ...:         false_values.append(val)
    ...: 
1000 loops, best of 3: 1.87 ms per loop

In [17]: %%timeit
    ...: res = {}
    ...: for val, bit in zip(value, selectors):
    ...:     res.setdefault(bit, []).append(val)
    ...: true_values, false_values = res.get(1, []), res.get(0, [])
    ...: 
100 loops, best of 3: 3.73 ms per loop

In [18]: from collections import defaultdict

In [19]: %%timeit
    ...: res = defaultdict(list)
    ...: for val, bit in zip(value, selectors):
    ...:     res[bit].append(val)
    ...: true_values, false_values = res.get(1, []), res.get(0, [])
    ...: 
100 loops, best of 3: 2.05 ms per loop
In [26]: %%timeit  # after conversion to numpy arrays
    ...: true_values = values[selectors == 0]
    ...: false_values = values[selectors == 1]
    ...: 
1000 loops, best of 3: 344 us per loop
In [31]: %%timeit
    ...: res = [[], []]
    ...: for val, bit in zip(value, selectors):
    ...:     res[bit].append(val)
    ...: true_values, false_values = res
    ...: 
100 loops, best of 3: 2.09 ms per loop
In [34]: from operator import not_

In [35]: %%timeit
    ...: true_values = list(compress(value, selectors))
    ...: false_values = list(compress(value, map(not_, selectors)))
    ...: 
1000 loops, best of 3: 1.44 ms per loop

显然numpy比其他的快得多，假设可以用numpy数组替换python列表。在

似乎itertools.compress是最快的非第三方解决方案，在1.44 ms。第二快的是朴素的for，在1.87处有一个if-else，其他的解决方案比{}稍微多一些。在

增加元素的数量我所看到的唯一变化是Jon Clement的defaultdict(list)解和newtower的[[], []]解变得比朴素的for+if-else（就像2%快500000）。compress仍然比其他人快30%，而{}仍然比compress快4倍左右。在

这种区别对你来说重要吗？如果不是（并配置文件以检查它是否是瓶颈！）我只考虑使用更具可读性的解决方案，这在很大程度上是主观的，由您决定。在

---

关于我得到的时间安排的最后一句话：

即使compress和您的双列表理解在列表上迭代两次，一个是最快的非第三方解决方案，另一个是最慢的。 这里您可以看到“python级循环”或“显式循环”与“C级循环”或“隐式循环”之间的区别。在

itertools.compress是在C中实现的，这使得它可以迭代而不需要太多的intepreter开销。正如你所看到的，这有很大的不同。在

您可以在numpy解决方案中看到更多这一点，它也执行两次迭代而不是一次迭代。在本例中，循环不仅是“在C级别”，而且还完全避免了调用pythonapi在数组上迭代，因为numpy有自己的C数据类型。在

这几乎是CPython中的一条规则：为了提高性能，尝试用隐式循环替换显式循环，使用内置函数或在C扩展中定义的函数。在

Guido van Rossum很清楚这一点，试着读他的An Optimization Anecdote。在

你可以在thisSO问题中找到另一个类似的例子（免责声明：接受的答案是我的。我利用了二分搜索和字符串相等（>；C级内置）来获得比纯python线性搜索更快的解决方案。在

Answer 2

您可以使用以下方法（或只使用collections.defaultdict(list)）：

res = {}
for val, bit in zip(vals, bits):
    res.setdefault(bit, []).append(val)
zeros, ones = res.get(0, []), res.get(1, [])

它只扫描一次列表，并对多个真/假值进行分组，但需要为新列表提供辅助存储。在

Answer 3

您可以尝试使用numpy数组来获得更好的性能：

els0 = vals[bits == 0]
els1 = vals[bits == 1]