如何用Pandas对DataFrame相邻同值states列分组求和？

原始数据 #

import pandas as pd

info = {
    'states': [-1, -1, -1, 1, 1, -1, 0, 1, 1, 1],
    'values': [34, 29, 28, 30, 35, 33, 33, 36, 40, 41]
}

df = pd.DataFrame(data=info)

原始DataFrame输出：

states  values
0       -1      34
1       -1      29
2       -1      28
3        1      30
4        1      35
5       -1      33
6        0      33
7        1      36
8        1      40
9        1      41

需求 #

仅对相邻的同值states行进行分组，对values列求和，得到如下结果：

states  values
0       -1      91  # 34+29+28
1        1      65  # 30+35
2       -1      33
3        0      33
4        1     117  # 36+40+41

解决方案 #

通过创建相邻同组的标识键，结合groupby实现求和：

# 生成分组标识：当当前行states与前一行不同时，累加生成新分组ID
df['group_id'] = (df['states'] != df['states'].shift()).cumsum()

# 按分组标识和states分组，对values求和
result = df.groupby(['group_id', 'states'], as_index=False)['values'].sum()

# 整理结果，移除group_id列并重置索引
result = result[['states', 'values']].reset_index(drop=True)

print(result)

输出结果：

states  values
0      -1      91
1       1      65
2      -1      33
3       0      33
4       1     117

代码说明 #

• df['states'].shift()：将states列向下偏移一行，用于对比当前行与前一行的states值
• (df['states'] != df['states'].shift()).cumsum()：生成分组ID，每遇到与前一行不同的states值时，ID递增，确保相邻同值的行属于同一分组
• groupby(['group_id', 'states'])：按分组ID和states分组，避免不同分组的相同states被合并
• 最后筛选需要的列并重置索引，得到目标格式的DataFrame

内容的提问来源于stack exchange，提问作者Federica F.