You can use itertools.compress
, which yields the elements that correspond to true in the selector.
However this would require to duplicate the bits
and invert a copy to select the elements for the zeros, which would end up with:
from operator import not_
true_values = list(compress(sequence, bits))
false_values = list(compress(sequence, map(not_, bits)))
I believe using a simple for
loop would be easier and faster, since it does a single iteration:
true_values = []
false_values = []
for bit, val in zip(bits, values):
if bit:
true_values.append(val)
else:
false_values.append(val)
Just as a curiosity, here are some micro benchmark with the various solutions:
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
Obviously numpy
is much faster then the rest, assuming that you can replace the python lists with numpy arrays.
It seems like itertools.compress
is the fastest non-3rd-party solution, at 1.44 ms
. The second fastest is the naive for
with an if-else
, at 1.87
, with other solutions taking slightly more than 2 ms
.
Increasing the number of elements the only changes I see is that the Jon Clement's defaultdict(list)
solution and newtower's [[], []]
solution become marginally faster than the naive for
+if-else
(like2%
faster at 500000
). compress
is still 30% faster than the others and numpy
is still about 4x faster than compress
.
Does this difference matter to you? If not(and profile to check whether it is a bottleneck!) I'd simply consider using the more readable solution, which is pretty much subjective and up to you.
A last remark on the timings I've obtained:
Even though both compress
and your double list-comprehension iterate over the list twice, one is the fastest non-3rd party solution, and the other is the slowest.
Here you can see the difference between a "python-level loop", or "explicit loop", versus a "C-level loop", or "implicit loop".
itertools.compress
is implemented in C
and this allows it to iterate without much of the intepreter overhead. And as you can see this makes a huge difference.
You can see this even more in the numpy
solution, which also performs two iterations instead of one. In this case, not only the loop is "at the C level", but it also completely avoid calling python APIs to iterate over the arrays, since numpy
has its own C data-types.
This is pretty much a rule in CPython: to improve performance try to replace explicit loops with implicit ones, using built-in functions or functions definined in C extensions.
Guido van Rossum is well-aware of this, try to read his An Optimization Anecdote.
you can find another such example in this SO question(disclaimer: the accepted answer is mine. I've exploited bisection search and string equality(-> the C-level built-in) to obtain a solution faster than a pure-python linear search).