如何用Pandas或Numpy按相邻值差值分组聚合原索引？

Pandas 版本 #

假设你有如下示例数据：

import numpy as np
import pandas as pd

# 原NumPy数组
arr = np.array([3, 1, 7, 15, 22, 25, 30])
# 原索引（默认从0开始的话，直接用np.arange生成即可）
original_indices = np.arange(len(arr))

步骤如下：

1. 将数据转为DataFrame，保留原索引

df = pd.DataFrame({'value': arr, 'original_idx': original_indices})

2. 按value列排序

df_sorted = df.sort_values('value')

3. 计算相邻值的差值，生成分组标签

# 相邻差值大于5则标记为新分组，cumsum生成连续的分组编号
group_labels = (df_sorted['value'].diff() > 5).cumsum()

4. 聚合原索引为字典或列表

• 转字典：键为分组编号，值为对应原索引列表

index_groups_dict = df_sorted.groupby(group_labels)['original_idx'].apply(list).to_dict()

• 转列表：每个元素为一组原索引

index_groups_list = df_sorted.groupby(group_labels)['original_idx'].apply(list).tolist()

运行后，index_groups_dict结果为{0: [1, 0, 2], 1: [3], 2: [4, 5], 3: [6]}，对应分组逻辑：排序后的值[1,3,7]差值均≤5归为一组；15与前一个值7差值>5单独成组；22、25差值≤5成组；30与25差值>5单独成组。

纯NumPy 等效实现 #

不需要Pandas的话，用纯NumPy操作即可完成：

import numpy as np

arr = np.array([3, 1, 7, 15, 22, 25, 30])
original_indices = np.arange(len(arr))

# 1. 获取排序后的原索引
sorted_indices = original_indices[np.argsort(arr)]
# 2. 获取排序后的数组
sorted_arr = arr[np.argsort(arr)]
# 3. 计算相邻差值，定位分组分割点
diff = np.diff(sorted_arr)
split_points = np.where(diff > 5)[0] + 1  # 分割点位置+1，因为diff是n-1个元素
# 4. 拆分原索引数组为分组
index_groups = np.split(sorted_indices, split_points)
# 转字典（可选）
index_groups_dict = {i: group.tolist() for i, group in enumerate(index_groups)}

运行后index_groups为数组列表：[array([1, 0, 2]), array([3]), array([4, 5]), array([6])]，转字典后与Pandas版本结果完全一致。

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