Question

我有一个名为vals的长度为N的浮点列表，以及一个名为bits的长度为N的0和1的列表。我想提取两个列表：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

您可以使用itertools.compress，它会在选择器中生成与true对应的元素。

然而，这需要复制bits并反转副本以选择零的元素，最终会得到：

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

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

true_values = []
false_values = []
for bit, val in zip(bits, values):
    if bit:
        true_values.append(val)
    else:
        false_values.append(val)

正如好奇心一样，这里有各种解决方案的微观基准：

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 if-else，1.87，其他解决方案略高于2 ms。

增加元素数量我看到的唯一变化是Jon Clement的defaultdict(list)解决方案和newtower的[[], []]解决方案比天真的for更快地略微 + if-else（在2%时更快500000）。 compress仍然比其他人快30％，numpy仍然比compress快4倍。

这种差异对你很重要吗？如果不是（并检查是否是瓶颈！）我只考虑使用更具可读性的解决方案，这是非常主观的，取决于你。

关于我获得的时间的最后一点评论：

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

itertools.compress在C中实现，这允许它在没有太多解释器开销的情况下进行迭代。正如你所看到的，这会产生巨大的差异。

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

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

Guido van Rossum非常清楚这一点，试着阅读他的An Optimization Anecdote。

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

Answer 2

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

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

Answer 3

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

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

它只扫描列表一次，并且分组超过true / false值，但确实需要新列表的辅助存储。

Answer 4

[val for idx, val in enumerate(values) if bits[idx]]

将为您提供与1匹配的值的子集。要获得包含两个子集（0和1）的列表，可以编写

true_vals, false_vals = [[val for idx, val in enumerate(values) if bits[idx]==bit] for bit in [0, 1]]

或

true_vals, false_vals = [[val[0] for val in zip(values, bits) if val[1]==bit] for bit in [0, 1]]

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

4 个答案: