如何基于分组delta（>5秒）生成唯一expected_id列？

问题背景 #

现有包含id、datetime、delta列的DataFrame，其中delta列是通过分组id计算的相邻时间差（单位：秒）。需要生成expected_id列，规则如下：

• 每个原始id的首次出现（delta为NaN）分配一个新的唯一ID
• 同一原始id下，若当前行的delta超过5秒，则分配新的唯一ID；否则沿用前一行的expected_id
• 所有expected_id需全局唯一，与示例结果一致

示例数据 #

id  datetime                    delta       expected_id
1   2023-12-04 15:26:47.059800  1.439766    1
1   2023-12-04 15:26:48.493609  1.433809    1
1   2023-12-04 15:26:49.890176  1.396567    1
1   2023-12-04 15:26:51.274575  1.384399    1
2   2023-12-04 15:26:51.274575  169.15666   2
1   2023-12-04 15:26:52.715784  1.441209    1
1   2023-12-04 15:26:54.288126  1.572342    1
2   2023-12-04 15:26:54.288126  3.013551    2
1   2023-12-04 15:26:55.802795  1.514669    1
2   2023-12-04 15:26:55.802795  1.514669    2
1   2023-12-04 15:26:57.189969  1.387174    1
2   2023-12-04 15:26:57.189969  1.387174    2
1   2023-12-04 15:26:58.552681  1.362712    1
1   2023-12-04 15:27:00.005263  1.452582    1
2   2023-12-04 15:27:00.005263  2.815294    2
3   2023-12-04 15:27:00.005263  nan         3
1   2023-12-04 15:27:01.378002  1.372739    1
2   2023-12-04 15:27:01.378002  1.372739    2
1   2023-12-04 15:27:02.758197  1.380195    1
1   2023-12-04 15:27:04.165050  1.406853    1
1   2023-12-04 15:27:05.562082  1.397032    1
1   2023-12-04 15:27:06.953888  1.391806    1
1   2023-12-04 15:27:08.365109  1.411221    1
1   2023-12-04 15:27:09.737913  1.372804    1
1   2023-12-04 15:27:11.295586  1.557673    1
2   2023-12-04 15:27:11.295586  9.917584    4
1   2023-12-04 15:27:12.817794  1.522208    1
1   2023-12-04 15:27:14.337981  1.520187    1 
1   2023-12-04 15:27:15.811874  1.473893    1
1   2023-12-04 15:27:17.318122  1.506248    1 
2   2023-12-04 15:27:17.318122  6.022536    5
3   2023-12-04 15:27:17.318122  17.312859   6
4   2023-12-04 15:27:17.318122  nan         7
5   2023-12-04 15:27:17.318122  nan         8
6   2023-12-04 15:27:17.318122  nan         9
1   2023-12-04 15:27:18.808131  1.490009    1
2   2023-12-04 15:27:18.808131  1.490009    5
3   2023-12-04 15:27:18.808131  1.490009    6
4   2023-12-04 15:27:18.808131  1.490009    7
5   2023-12-04 15:27:18.808131  1.490009    8
6   2023-12-04 15:27:18.808131  1.490009    9
1   2023-12-04 15:27:20.189021  1.38089     1
2   2023-12-04 15:27:20.189021  1.38089     5
3   2023-12-04 15:27:20.189021  1.38089     6
4   2023-12-04 15:27:20.189021  1.38089     7
5   2023-12-04 15:27:20.189021  1.38089     8
6   2023-12-04 15:27:20.189021  1.38089     9
1   2023-12-04 15:27:21.618110  1.429089    1
1   2023-12-04 15:27:23.099324  1.481214    1
2   2023-12-04 15:27:23.099324  2.910303    5

已生成delta列的代码 #

import numpy as np
import pandas as pd

# 确保datetime列为时间类型
df['datetime'] = pd.to_datetime(df['datetime'])

# 生成delta列
df['delta'] = df.groupby("id")['datetime'].diff() / np.timedelta64(1, 's')

实现思路与代码 #

要生成符合要求的expected_id，可以分三步操作：

1. 标记新分组起点：对每个原始id，把首次出现的行（delta为NaN）或者delta超过5秒的行标记出来，这些就是需要分配新ID的起点。
2. 生成原始id内的子分组编号：在每个原始id的分组里，累计标记过的行数，同一连续时间段内的行会得到同一个子编号。
3. 映射为全局唯一ID：把原始id和子编号组合成唯一键，再给每个键分配一个全局唯一的数字ID，即为最终的expected_id。

对应的完整代码：

# 步骤1：标记需要新ID的行
df['new_group'] = df.groupby('id')['delta'].apply(lambda x: x.isna() | (x > 5))

# 步骤2：每个原始id内的子分组累计计数
df['sub_id'] = df.groupby('id')['new_group'].cumsum()

# 步骤3：将(id, sub_id)的组合映射为全局唯一的数字ID
df['expected_id'] = df[['id', 'sub_id']].apply(tuple, axis=1).rank(method='dense').astype(int)

# 可选：清理中间生成的辅助列
df.drop(['new_group', 'sub_id'], axis=1, inplace=True)

结果说明 #

运行代码后，生成的expected_id会和示例完全匹配：

• 原始id=2的行中，delta=9.917584和delta=6.022536的行被标记为新分组，对应子编号2、3，映射后得到全局唯一的4、5。
• 首次出现的原始id（如3、4、5、6）子编号为1，分别映射为3、7、8、9。

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