如何用Pandas生成目标累计求和的DataFrame？

输入数据 #

data = {'count': [2,2,2,2,1,0,0,1,1]}

期望输出 #

data = {'count': [2,2,2,2,3,3,3,4,4]}

问题分析 #

原代码的累加逻辑未按连续非零值的分组进行递增处理，无法实现目标效果。

解决方案代码 #

import pandas as pd

data = {'count': [2,2,2,2,1,0,0,1,1]}
df = pd.DataFrame(data)

# 用前向填充替换0值
df['count'] = df['count'].replace(0, method='ffill')

# 标记连续相同值的分组：当前值与前一个值不同时，分组ID+1
df['group'] = (df['count'] != df['count'].shift()).cumsum()

# 提取每个分组的初始值，计算分组对应的偏移量
group_initial = df.groupby('group')['count'].first()
offset = group_initial.cumsum() - group_initial

# 将偏移量映射到原数据，得到最终结果
df['count'] += df['group'].map(offset)

# 删除临时生成的group列
df = df.drop('group', axis=1)

print(df)

代码说明 #

1. 填充0值：把序列中的0替换为前一个非零值，得到中间结果[2,2,2,2,1,1,1,1,1]。
2. 分组标记：通过对比当前值与前一个值的差异，生成分组ID，把连续相同的数值划分为同一组，分组结果为[1,1,1,1,2,2,2,3,3]。
3. 计算偏移量：对每个分组的初始值做累加运算，再减去初始值得到该分组需要叠加的偏移量。这里初始值为[2,1,1]，累加后为[2,3,4]，对应偏移量就是[0,2,3]。
4. 叠加偏移量：将每个分组对应的偏移量加到原数值上，最终得到目标序列。

内容的提问来源于stack exchange，提问作者Magic Stick