I have a (long) list in which zeros and ones appear at random:
list_a = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]
I want to get the list_b
- sum of the list up to where 0 appears
where 0 appears, retain 0 in the list
list_b = [1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
I can implement this as follows:
list_b = []
for i, x in enumerate(list_a):
if x == 0:
list_b.append(x)
else:
sum_value = 0
for j in list_a[i::-1]:
if j != 0:
sum_value += j
else:
break
list_b.append(sum_value)
print(list_b)
but the actual list's length is very long.
So, I want to improve code for high speed. (if it is not readable)
I change the code like this:
from itertools import takewhile
list_c = [sum(takewhile(lambda x: x != 0, list_a[i::-1])) for i, d in enumerate(list_a)]
print(list_c)
But it is not fast enough. How can I do it in more efficient way?
You're overthinking this.
Option 1
You can just iterate over the indices and update accordingly (computing the cumulative sum), based on whether the current value is 0 or not.
data = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]
for i in range(1, len(data)):
if data[i]:
data[i] += data[i - 1]
That is, if the current element is non-zero, then update the element at the current index as the sum of the current value, plus the value at the previous index.
print(data)
[1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
Note that this updates your list in place. You can create a copy in advance if you don't want that - new_data = data.copy() and iterate over new_data in the same manner.
Option 2
You can use the pandas API if you need performance. Find groups based on the placement of 0s, and use groupby + cumsum to compute group-wise cumulative sums, similar to above:
import pandas as pd
s = pd.Series(data)
data = s.groupby(s.eq(0).cumsum()).cumsum().tolist()
print(data)
[1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
Performance
First, the setup -
data = data * 100000
s = pd.Series(data)
Next,
%%timeit
new_data = data.copy()
for i in range(1, len(data)):
if new_data[i]:
new_data[i] += new_data[i - 1]
328 ms ± 4.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
And, timing the copy separately,
%timeit data.copy()
8.49 ms ± 17.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
So, the copy doesn't really take much time. Finally,
%timeit s.groupby(s.eq(0).cumsum()).cumsum().tolist()
122 ms ± 1.69 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
The pandas approach is conceptually linear (just like the other approaches) but faster by a constant degree because of the implementation of the library.
If you want a compact native Python solution that is probably the most memory efficient, although not the fastest (see the comments), you could draw extensively from itertools:
>>> from itertools import groupby, accumulate, chain
>>> list(chain.from_iterable(accumulate(g) for _, g in groupby(list_a, bool)))
[1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
The steps here are: group the list into sublists based on presence of 0 (which is falsy), take the cumulative sum of the values within each sublist, flatten the sublists.
As Stefan Pochmann comments, if your list is binary in contents (like consisting of only 1s and 0s only) then you don't need to pass a key to groupby() at all and it will fall back on the identity function. This is ~30% faster than using bool for this case:
>>> list(chain.from_iterable(accumulate(g) for _, g in groupby(list_a)))
[1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
Personally I would prefer a simple generator like this:
def gen(lst):
cumulative = 0
for item in lst:
if item:
cumulative += item
else:
cumulative = 0
yield cumulative
Nothing magic (when you know how yield works), easy to read and should be rather fast.
If you need more performance you could even wrap this as Cython extension type (I'm using IPython here). Thereby you lose the "easy to understand" portion and it's requiring "heavy dependencies":
%load_ext cython
%%cython
cdef class Cumulative(object):
cdef object it
cdef object cumulative
def __init__(self, it):
self.it = iter(it)
self.cumulative = 0
def __iter__(self):
return self
def __next__(self):
cdef object nxt = next(self.it)
if nxt:
self.cumulative += nxt
else:
self.cumulative = 0
return self.cumulative
Both need to be consumed, for example using list to give the desired output:
>>> list_a = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]
>>> list(gen(list_a))
[1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
>>> list(Cumulative(list_a))
[1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
However since you asked about speed I wanted to share the results from my timings:
import pandas as pd
import numpy as np
import random
import pandas as pd
from itertools import takewhile
from itertools import groupby, accumulate, chain
def MSeifert(lst):
return list(MSeifert_inner(lst))
def MSeifert_inner(lst):
cumulative = 0
for item in lst:
if item:
cumulative += item
else:
cumulative = 0
yield cumulative
def MSeifert2(lst):
return list(Cumulative(lst))
def original1(list_a):
list_b = []
for i, x in enumerate(list_a):
if x == 0:
list_b.append(x)
else:
sum_value = 0
for j in list_a[i::-1]:
if j != 0:
sum_value += j
else:
break
list_b.append(sum_value)
def original2(list_a):
return [sum(takewhile(lambda x: x != 0, list_a[i::-1])) for i, d in enumerate(list_a)]
def Coldspeed1(data):
data = data.copy()
for i in range(1, len(data)):
if data[i]:
data[i] += data[i - 1]
return data
def Coldspeed2(data):
s = pd.Series(data)
return s.groupby(s.eq(0).cumsum()).cumsum().tolist()
def Chris_Rands(list_a):
return list(chain.from_iterable(accumulate(g) for _, g in groupby(list_a, bool)))
def EvKounis(list_a):
cum_sum = 0
list_b = []
for item in list_a:
if not item: # if our item is 0
cum_sum = 0 # the cumulative sum is reset (set back to 0)
else:
cum_sum += item # otherwise it sums further
list_b.append(cum_sum) # and no matter what it gets appended to the result
def schumich(list_a):
list_b = []
s = 0
for a in list_a:
s = a+s if a !=0 else 0
list_b.append(s)
return list_b
def jbch(seq):
return list(jbch_inner(seq))
def jbch_inner(seq):
s = 0
for n in seq:
s = 0 if n == 0 else s + n
yield s
# Timing setup
timings = {MSeifert: [],
MSeifert2: [],
original1: [],
original2: [],
Coldspeed1: [],
Coldspeed2: [],
Chris_Rands: [],
EvKounis: [],
schumich: [],
jbch: []}
sizes = [2**i for i in range(1, 20, 2)]
# Timing
for size in sizes:
print(size)
func_input = [int(random.random() < 0.75) for _ in range(size)]
for func in timings:
if size > 10000 and (func is original1 or func is original2):
continue
res = %timeit -o func(func_input) # if you use IPython, otherwise use the "timeit" module
timings[func].append(res)
%matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure(1)
ax = plt.subplot(111)
baseline = MSeifert2 # choose one function as baseline
for func in timings:
ax.plot(sizes[:len(timings[func])],
[time.best / ref.best for time, ref in zip(timings[func], timings[baseline])],
label=func.__name__) # you could also use "func.__name__" here instead
ax.set_ylim(0.8, 1e4)
ax.set_yscale('log')
ax.set_xscale('log')
ax.set_xlabel('size')
ax.set_ylabel('time relative to {}'.format(baseline)) # you could also use "func.__name__" here instead
ax.grid(which='both')
ax.legend()
plt.tight_layout()
In case you're interested in the exact results I put them in this gist.
It's a log-log plot and relative to the Cython answer. In short: The lower the faster and the range between two major tick represents one order of magnitude.
So all solutions tend to be within one order of magnitude (at least when the list is big) except for the solutions you had. Strangely the pandas solution is quite slow compared to the pure Python approaches. However the Cython solution beats all of the other approaches by a factor of 2.
You are playing with the indices too much in the code you posted when you do not really have to. You can just keep track of a cumulative sum and reset it to 0 every time you meet a 0.
list_a = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]
cum_sum = 0
list_b = []
for item in list_a:
if not item: # if our item is 0
cum_sum = 0 # the cumulative sum is reset (set back to 0)
else:
cum_sum += item # otherwise it sums further
list_b.append(cum_sum) # and no matter what it gets appended to the result
print(list_b) # -> [1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
It doesn't have to be as complicated as made in the question asked, a very simple approach could be this.
list_a = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]
list_b = []
s = 0
for a in list_a:
s = a+s if a !=0 else 0
list_b.append(s)
print list_b
I would use a generator if you want performance (and it's simple too).
def weird_cumulative_sum(seq):
s = 0
for n in seq:
s = 0 if n == 0 else s + n
yield s
list_b = list(weird_cumulative_sum(list_a_))
I don't think you'll get better than that, in any case you'll have to iterate over list_a at least once.
Note that I called list() on the result to get a list like in your code but if the code using list_b is iterating over it only once with a for loop or something there is no use converting the result to a list, just pass it the generator.
Starting Python 3.8, and the introduction of assignment expressions (PEP 572) (:= operator), we can use and increment a variable within a list comprehension:
# items = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]
total = 0
[total := (total + x if x else x) for x in items]
# [1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
This:
- Initializes a variable
totalto0which symbolizes the running sum - For each item, this both:
- either increments
totalwith the current looped item (total := total + x) via an assignment expression or set it back to0if the item is0 - and at the same time, maps
xto the new value oftotal
- either increments
来源：https://stackoverflow.com/questions/48805035/compute-the-cumulative-sum-of-a-list-until-a-zero-appears