Pandas中连续相同A值分组的部分和计算问题

你说得没错，普通的groupby('A')确实会把所有相同A值的行归为一组，不管它们在DataFrame里是不是连续的，这就没法满足你要按连续相邻分组的需求啦。咱们可以通过先创建一个「连续分组标识」的方式来解决这个问题，具体步骤如下：

1. 先构造你的示例DataFrame #

首先咱们把你给出的数据转换成可运行的代码：

import pandas as pd

df = pd.DataFrame({
    'A': [1.2, 1.2, 1.2, 2.3, 2.3, 1.2],
    'B': [1, 6, 4, 2, 5, 7]
})

2. 创建连续分组的标识键 #

核心思路是：当某一行的A值和上一行不同时，就标记为一个新的分组。我们可以用ne()（不等于）结合shift()（上移一行）来判断，再用cumsum()（累加）生成分组编号：

# 生成连续分组的编号，相同连续A值会得到同一个编号
group_key = df['A'].ne(df['A'].shift()).cumsum()

运行后group_key的结果是[1,1,1,2,2,3]，完美把连续的相同A值分成了3个组。

3. 按新的分组键求和并整理结果 #

接下来用这个分组键进行聚合，同时保留每个连续组的起始索引（和你预期的输出索引一致），并对B列求和：

# 分组聚合：取每个组的A值（因为同一组A值相同，取第一个即可），对B列求和
result = df.groupby(group_key).agg(
    A=('A', 'first'),
    B=('B', 'sum')
)

# 把结果的索引替换成每个连续组的起始索引
result = result.set_index(df.groupby(group_key).apply(lambda x: x.index[0]))

最终结果 #

运行后得到的result就是你想要的：

A   B
0  1.2  11
3  2.3   7
5  1.2   7

这样就完美实现了按连续相邻的相同A值分组求和的需求啦～

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