如何为Pandas Series中的连续数据区域分配唯一ID？

需求说明 #

需要将Pandas Series中满足相邻元素为相等或递增1的连续数据区域划分为同一组，并为每个组分配唯一ID。

示例输入：

import pandas as pd
df = pd.Series(data = [0,1,1,1,2,2,2,16,17,18,18,18,18,1,1,1,1,2,2,200,201,3,4])

期望输出：

pd.Series(data = [0,0,0,0,0,0,0,1,1,1,1,1,1,2,2,2,2,2,2,3,3,4,4])

你尝试的代码无法实现需求（还存在语法错误）：

df= pd.concat([df, (df==0).cumsum()], axis = 1, keys = ['val', 'flag']

解决方案 #

核心思路是：先判断每个位置是否是新分组的起始点（即当前元素与前一个元素的差值既不是0也不是1），然后对这些起始点的标记进行累加，得到分组ID。

具体代码：

import pandas as pd

# 示例数据
s = pd.Series(data = [0,1,1,1,2,2,2,16,17,18,18,18,18,1,1,1,1,2,2,200,201,3,4])

# 计算相邻元素的差值
diff = s.diff()

# 标记新分组的起始点：差值不在{0,1}范围内，或者是第一个元素（diff为NaN）
new_group = (diff.isna()) | (~diff.isin([0, 1]))

# 对新分组标记累加，得到从0开始的分组ID
group_id = new_group.cumsum() - 1

print(group_id)

运行结果与期望输出完全一致：

0     0
1     0
2     0
3     0
4     0
5     0
6     0
7     1
8     1
9     1
10    1
11    1
12    1
13    2
14    2
15    2
16    2
17    2
18    2
19    3
20    3
21    4
22    4
dtype: int64

代码解释 #

1. s.diff()：计算当前元素与前一个元素的差值，第一个元素因无前驱值，差值为NaN。
2. new_group：标记新分组的起点，包含两种情况：一是第一个元素，二是当前元素与前一个元素的差值既不是0（元素未变化）也不是1（元素递增1）的位置。
3. new_group.cumsum() - 1：对标记为True的起始点进行累加，得到从1开始的分组序号，减1后转换为从0开始的唯一ID。

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