Pandas按a1列相邻值预设差值规则实现行索引分组

核心思路 #

已知a1列已提前按顺序排列，只需三步完成分组：

1. 计算a1列相邻行的差值，标记差值不等于预设阈值的位置为新分组起点
2. 对分组起点标记做累计求和，为每一行生成对应的连续组号
3. 按组号聚合行索引，整理为要求的嵌套列表格式

完整可运行代码 #

import pandas as pd
import numpy as np

# 构造示例数据表
data = [
    [5, 2],
    [100, 23],
    [101, -2],
    [303, 9],
    [304, 4],
    [709, 14],
    [710, 3],
    [711, 3],
    [988, 21]
]
columns = ['a1', 'a2']
df = pd.DataFrame(data=data, columns=columns)

# 配置分组阈值
threshold = 1
# 标记新分组起点：相邻差值不等于阈值时为新组起点，首行默认作为第一个组起点
is_new_group = df['a1'].diff() != threshold
# 累计求和生成组号
df['group_tag'] = is_new_group.cumsum()
# 聚合得到分组索引列表
group_result = df.groupby('group_tag').apply(lambda x: x.index.tolist()).tolist()

# 验证输出
print(group_result)
# 运行结果：[[0], [1, 2], [3, 4], [5, 6, 7], [8]]

逻辑说明 #

• 差值计算使用Series.diff()方法，第一行无前置值时返回NaN，和阈值比较结果为True，天然作为第一个分组的起点，无需额外填充处理。
• 布尔序列累计求和的逻辑：每遇到一个True（新分组起点），累计值加1，同组内所有行的累计值保持一致，即可实现连续分段的组号标记。
• 如果需要调整分组判定阈值，仅需修改threshold变量的取值即可，核心逻辑无需改动。

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