如何基于Pandas实现按连续Role及索引差≤5分组？

原始数据集 #

原有分组代码及结果 #

此前使用以下代码实现仅按连续相同Role字段分组：

df = df.groupby((df.Role!= df.Role.shift()).cumsum()).agg(
            Role = ('Role', 'first'),
            Name = ('Name', ' '.join), 
            Grade = ('Grade', 'mean')
        ).reset_index(drop=True)

执行后得到结果：

新增规则与预期结果 #

现在需要新增分组条件：仅当当前行索引与前一行索引的差值≤5时，才对连续相同Role的行进行分组，预期结果如下：

解决方案 #

核心思路 #

同时满足两个分组中断条件时开启新分组：

• 当前行Role与前一行不同
• 当前行索引与前一行索引差值>5
  只要触发其中一个条件，就生成新的分组标识。

实现代码 #

# 计算分组标识，同时判断Role变化和索引差值
df['group_id'] = ((df['Role'] != df['Role'].shift()) | (df['索引'] - df['索引'].shift() > 5)).cumsum()

# 按分组标识聚合，保留两位小数匹配预期结果
result_df = df.groupby('group_id').agg(
    Role=('Role', 'first'),
    Name=('Name', ' '.join),
    Grade=('Grade', lambda x: round(x.mean(), 2))
).reset_index(drop=True)

# 重置结果索引为从1开始
result_df.index = result_df.index + 1

代码说明 #

1. 用shift()获取前一行的Role和索引值，生成布尔序列标记分组中断点：
   • df['Role'] != df['Role'].shift()：Role变化时触发分组中断
   • df['索引'] - df['索引'].shift() > 5：索引差值超5时触发分组中断
   • 用|合并两个条件，cumsum()将布尔序列转为连续的分组ID
2. 按分组ID执行聚合操作，和原有逻辑一致，新增小数位数处理匹配预期格式
3. 最后重置结果索引为从1开始，对齐预期输出的索引格式

高效性说明 #

该方法基于pandas矢量化操作，时间复杂度为O(n)，无需循环遍历，在大数据集上也能保持高效。若处理超大规模数据，可提前确保索引列是数值类型，避免类型转换的额外开销。

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