如何为含时间戳索引的Pandas DataFrame高效分配测量ID？

问题描述 #

有一个包含400组独立测量数据的Pandas DataFrame，每组是约1000行的时间序列，总数据量约40万行，索引为时间戳。需要给每组测量分配唯一ID作为新列，方便通过df[df["ID"] == certain_id]快速查询。

当前采用逐行遍历对比时间戳的方式实现，但耗时极长，原始代码如下：

t_0 = df.index[0] # very first timestamp
id = 1 # first ID

df.loc[df.index[0], "ID"] = 1

for row in df_raw.itertuples():
   try:
      t_1 = row.Index
   except IndexError:
      break
   if t_1 >= t_0 + pd.to_timedelta(min_diff): # Check for threshold
      df.loc[row.Index, "ID"] = id
      id += 1
      t_0 = t_1
      continue
   else:
      df.loc[row.Index, "ID"] = id
      t_0 = t_1

优化方案 #

Pandas的逐行遍历+loc赋值效率极低，尤其是大数据集场景，改用向量化操作可以将耗时降低几个数量级，核心思路是利用时间差计算和累加生成ID：

1. 计算相邻索引的时间差：用df.index.to_series().diff()得到每行与上一行的时间间隔
2. 标记超过阈值的分界点：判断时间差是否大于设定的min_diff，得到布尔数组
3. 生成分组ID：对分界点数组做累加（cumsum()），再加上1（因为初始ID从1开始）

具体代码 #

import pandas as pd

# 计算相邻时间戳的差值
time_diff = df.index.to_series().diff()

# 标记需要新ID的位置（第一行默认是新分组，所以fillna(True)）
new_group = time_diff > pd.to_timedelta(min_diff)
new_group = new_group.fillna(True)

# 生成ID列：累加标记，从1开始
df["ID"] = new_group.cumsum()

代码解释 #

• df.index.to_series().diff()：将索引转为Series后计算差值，得到每行与前一行的时间间隔，第一行差值为NaN
• new_group.fillna(True)：第一行没有前序数据，默认标记为新分组
• new_group.cumsum()：对布尔数组累加，每遇到一个True就会让累加值加1，正好对应分组ID的递增

这种方法全程用Pandas内置的向量化函数，避免了循环和逐行赋值，40万行数据的处理时间会从分钟级降到毫秒级。

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