如何加速Python中DataFrame列的非零值分段求和代码？

现有一个DataFrame的df.Ah列，值为正数或零。需要将零之前的所有非零值求和，再紧跟一个零，结果存入新的DataFrame。例如：

• 输入：df.Ah = [1,2,3,0,46,0,24,1]
• 输出：[6,0,46,0,25]

当前尝试的代码如下，处理3100万条数据耗时约5小时，效率极低：

lst = df.Ah;
lst1 = pd.DataFrame(index=np.arange(100000000)); #if I replace values of a pre-existing column code #speeds up

summ = 0;
for i , elem in enumerate(lst):
     
    if elem != 0:
        summ = summ + elem;
    else:
        if summ:
            lst1.loc[i,'Ah']=sum;
            
        lst1.loc[i,'Ah']=elem;
        summ = 0;
...

if summ:
    lst1.iloc[i+1, 0] = summ;

直接使用pandas的向量化分组操作，完全避免Python循环，可将处理时间压缩到分钟级：

核心思路 #

通过生成分组标签，把连续非零值与后续的零归为同一组，再对每组求和，最后拼接出符合要求的结果序列。

代码实现 #

import pandas as pd
import numpy as np

# 替换成你的df['Ah']
s = pd.Series([1,2,3,0,46,0,24,1])

# 生成分组ID：遇到0时分组ID递增，把连续非零值和后续零划为一组
group_ids = (s == 0).cumsum()

# 计算每个分组的总和
group_sums = s.groupby(group_ids).sum()

# 构造结果列表
result = []
for idx, sum_val in group_sums.items():
    result.append(sum_val)
    # 若当前分组包含0，则追加一个0
    if (s[group_ids == idx] == 0).any():
        result.append(0)

# 转为目标DataFrame
result_df = pd.DataFrame({'Ah': result})

效率提升原因 #

• 摒弃Python循环，所有核心操作由pandas底层C语言实现，千万级数据处理速度大幅提升。
• 无需提前创建超大空DataFrame，直接构造结果序列，减少内存占用和无效赋值操作。

内容的提问来源于stack exchange，提问作者Pietro g