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

Question

我有一個長度為N的float列表，稱為vals ，長度為N的0和1s列表，稱為bits 。 我想提取兩個列表： vals的元素，它們對應於bits的0，其余元素對應於bits的1s。 我目前正在做：

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數組替換python列表，那么numpy的速度要比其余的要快得多。

看來itertools.compress是最快的非第三方解決方案，速度為1.44 ms 。 第二個最快的是天真for用if-else ，在1.87 ，與其它解決方案略服用超過2 ms 。

增加元件的數量我看到的唯一變化是，喬恩·克萊門特的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）對此非常了解，請嘗試閱讀他的《優化軼事》 。

您可以在此 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 s匹配的值的子集。 要獲得包含兩個子集（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]]