python - 计算列表的累积和，直到出现零

Question

我有一个(长)列表，其中随机出现零和一:

list_a = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]

我要获取列表_b

• 列表的总和，直到出现 0 的位置

• 出现0的地方，在列表中保留0

  list_b = [1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]
  

我可以这样实现:

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)

但实际列表的长度很长。

所以，我想改进代码以提高速度。 (如果它不可读)

我把代码改成这样:

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)

但这还不够快。我怎样才能更有效地做到这一点？

Answer 1

你想多了。

选项 1
您可以根据当前值是否为 0 来迭代索引并相应地更新(计算累积总和)。

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] 

即如果当前元素不为零，则将当前索引处的元素更新为当前值与前一个索引处的值之和。

print(data)
[1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]

请注意，这会更新您的列表。如果您不想这样做，可以提前创建一个副本 - new_data = data.copy() 并以相同的方式迭代 new_data。

---

选项 2
如果您需要性能，可以使用 pandas API。根据 0 的位置查找分组，并使用 groupby + cumsum 计算分组累积和，类似上面:

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]

---

性能

首先，设置-

data = data * 100000
s = pd.Series(data)

接下来，

%%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)

并且，单独定时复制，

%timeit data.copy()
8.49 ms ± 17.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

所以，复制并不需要太多时间。最后，

%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)

pandas 方法在概念上是线性的(就像其他方法一样)，但由于库的实现，速度更快。