如何为数据集中连续存在（中断≤5秒）的组分配唯一累计ID？

我有一个包含group_id的数据集，需要为持续存在的组分配唯一ID，规则是：

• 若组消失时长不超过5秒，保留原累计ID；
• 若组消失时长超过5秒，分配新的累计ID。

数据集如下：

group_id    dt2         expect_id
0   nan 2023-11-28 17:43:09.900628  1
1   1   2023-11-28 17:43:11.322793  2
2   1   2023-11-28 17:43:12.660818  2
3   1   2023-11-28 17:43:14.119043  2
4   1   2023-11-28 17:43:15.550513  2
5   2   2023-11-28 17:43:15.550513  3
6   3   2023-11-28 17:43:15.550513  4
7   4   2023-11-28 17:43:15.550513  5
8   1   2023-11-28 17:43:16.973557  2
9   2   2023-11-28 17:43:16.973557  3
10  3   2023-11-28 17:43:16.973557  4
11  4   2023-11-28 17:43:16.973557  5
12  1   2023-11-28 17:43:18.335619  2
13  2   2023-11-28 17:43:18.335619  3
14  3   2023-11-28 17:43:18.335619  4
15  4   2023-11-28 17:43:18.335619  5
16  1   2023-11-28 17:43:19.738230  2
17  2   2023-11-28 17:43:19.738230  3
18  3   2023-11-28 17:43:19.738230  4
19  4   2023-11-28 17:43:19.738230  5
20  1   2023-11-28 17:43:21.110693  2
21  2   2023-11-28 17:43:21.110693  3
22  1   2023-11-28 17:43:22.571257  2
23  2   2023-11-28 17:43:22.571257  3
24  1   2023-11-28 17:43:24.000589  2
25  1   2023-11-28 17:43:25.429940  2
26  2   2023-11-28 17:43:25.429940  3
27  1   2023-11-28 17:43:26.851142  2
28  2   2023-11-28 17:43:26.851142  3
29  1   2023-11-28 17:43:28.256274  2
30  nan 2023-11-28 17:43:29.617541  4
31  nan 2023-11-28 17:43:30.974490  4
32  nan 2023-11-28 17:43:32.360739  4
33  1   2023-11-28 17:43:33.730457  5
34  2   2023-11-28 17:43:33.730457  6
35  1   2023-11-28 17:43:35.270380  5
34  2   2023-11-28 17:43:35.270380  6
35  1   2023-11-28 17:43:35.270380  5
34  2   2023-11-28 17:43:36.730457  6
35  1   2023-11-28 17:43:36.270380  5

我尝试用cumsum()创建unique_id列，但这个方法只会根据group_id的变化分配新ID，没考虑时间间隔：

df['group_id'] = (df['group_id'].eq(0) | (df['group_id'] != df['group_id'].shift())).cumsum()

请问怎么实现仅当group_id消失超过5秒时才分配新累计ID的逻辑？

要实现这个逻辑，核心是按group_id分组，跟踪每组上次出现的时间，判断间隔是否超过5秒，再决定要不要分配新ID。具体步骤如下：

1. 先把时间列转成可计算的格式 #

首先得确保dt2是datetime类型，不然没法算时间差：

import pandas as pd

df['dt2'] = pd.to_datetime(df['dt2'])

2. 给每个group_id找上次出现的时间 #

过滤掉空的group_id，按组分组后，用shift()拿到每组上一次出现的时间，再计算时间差（单位秒）：

# 先处理非空的group_id
df_non_nan = df.dropna(subset=['group_id']).copy()
# 按group_id分组，取上一行的时间
df_non_nan['prev_dt'] = df_non_nan.groupby('group_id')['dt2'].shift(1)
# 计算当前时间和上次出现的时间差（秒）
df_non_nan['time_diff'] = (df_non_nan['dt2'] - df_non_nan['prev_dt']).dt.total_seconds()

3. 标记哪些情况需要新ID #

对每个组来说，第一次出现肯定要新ID；如果之后再出现时，和上次的间隔超过5秒，也得给新ID：

# 第一次出现的行标记为需要新ID（1）
df_non_nan['new_id_flag'] = df_non_nan.groupby('group_id')['prev_dt'].isna().astype(int)
# 时间间隔超过5秒的行也标记为1
df_non_nan.loc[df_non_nan['time_diff'] > 5, 'new_id_flag'] = 1
# 按组累计标记，得到每个组内部的ID序号
df_non_nan['unique_id'] = df_non_nan.groupby('group_id')['new_id_flag'].cumsum()

4. 处理空的group_id #

看你的expect_id，空组是连续的，直接给同一个ID就行：

df_nan = df[df['group_id'].isna()].copy()
df_nan['unique_id'] = 4  # 和你的expect_id对应

5. 合并结果并生成全局ID #

把两部分数据合并，再把每个(group_id, unique_id)组合映射成全局的数字ID，就和你的expect_id一致了：

# 合并数据并按原顺序排序
result = pd.concat([df_non_nan, df_nan]).sort_index()
# 给每个(group_id, unique_id)分配全局唯一数字
id_mapping = result.dropna(subset=['group_id']).drop_duplicates(['group_id', 'unique_id'])
id_mapping['global_id'] = range(2, len(id_mapping)+2)  # 从2开始对应示例里的ID
# 把映射合并回去
result = result.merge(id_mapping[['group_id', 'unique_id', 'global_id']], on=['group_id', 'unique_id'], how='left')
# 第一行的空组单独设为1
result.loc[result.index == 0, 'global_id'] = 1

简化版代码 #

如果不想拆分处理空组，也可以用更简洁的写法：

import pandas as pd

df['dt2'] = pd.to_datetime(df['dt2'])

# 标记非空group_id的行
mask = df['group_id'].notna()
# 计算每个非空group_id上次出现的时间和间隔
df.loc[mask, 'prev_dt'] = df.loc[mask].groupby('group_id')['dt2'].shift()
df['time_diff'] = (df['dt2'] - df['prev_dt']).dt.total_seconds()

# 标记需要新ID的情况：第一次出现 或 间隔超5秒
df['new_id'] = 0
df.loc[mask & df['prev_dt'].isna(), 'new_id'] = 1
df.loc[mask & (df['time_diff'] > 5), 'new_id'] = 1

# 计算每个组内部的ID序号
df['local_id'] = df.groupby('group_id')['new_id'].cumsum()

# 生成全局唯一标识，再转成数字ID
df['global_tag'] = df.apply(lambda x: f"{x['group_id']}_{x['local_id']}" if pd.notna(x['group_id']) else 'nan_group', axis=1)
# 把标签映射成数字，和expect_id对应
tag_to_id = {tag:i+1 for i, tag in enumerate(df['global_tag'].unique())}
df['expect_id'] = df['global_tag'].map(tag_to_id)

这样跑出来的expect_id就和你想要的结果完全匹配了。

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