Pandas DataFrame按no列条件计算分组差值及总和的实现问题

核心思路 #

• 先为DataFrame中的连续相同no值打上唯一的区块ID，区分不同的连续值段
• 提取所有no=1区块的首行count作为value1
• 匹配每个no=1区块后第一个非1区块的首行count作为value2
• 计算差值后汇总求和

完整代码 #

import pandas as pd

# 构造测试数据
data = [[1, 10], [1, 5], [1, 4], [0, 3], [0, 3], [1, 2], [1, 1], [1, 5], [2, 8], [4, 9]]
df = pd.DataFrame(data, columns = ['no', 'count'])

# 标记连续相同no的区块
df['block_id'] = (df['no'] != df['no'].shift()).cumsum()

# 提取所有no=1区块的首行count值
block_1 = df[df['no'] == 1].groupby('block_id')['count'].first().reset_index()
block_1.columns = ['block_id', 'value1']

# 提取所有非1区块的首行count值
block_not1 = df[df['no'] != 1].groupby('block_id')['count'].first().reset_index()
block_not1.columns = ['block_id', 'value2']

# 匹配计算差值
diffs = []
for idx, row in block_1.iterrows():
    # 取当前1区块之后第一个非1区块的value2
    next_not1 = block_not1[block_not1['block_id'] > row['block_id']].iloc[0]
    # 按预期输出调整计算逻辑，如严格按需求写value2-value1可交换前后顺序
    diff = row['value1'] - next_not1['value2']
    diffs.append(diff)

print("所有差值：", diffs)
print("差值总和：", sum(diffs))

输出结果 #

所有差值： [7, -6]
差值总和： 1

补充说明 #

如果严格按照需求描述的value2 - value1计算，仅需修改差值计算行为diff = next_not1['value2'] - row['value1']，运行后得到差值为[-7, 6]，总和为-1，与预期符号相反，属于需求描述笔误。

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