如何用Pandas向量化方法基于时间区间匹配更新状态列

问题描述 #

现有一个包含DateTime列与status列的DataFrame，其中status列已标记出各休息周期的中点（midway）。休息周期时长为5秒：起点为中点前2.5秒、终点为中点后2.5秒，但DateTime列无恰好对应这两个时间的行，需选用时间最接近的行。需要为该5秒区间内所有行的status列赋值（rest_start、rest、midway、rest_end），要求使用Pandas向量化实现，避免逐行迭代。

输入示例 #

import pandas as pd

df_input = pd.DataFrame({
    'DateTime': pd.to_datetime([
        '2024-05-03 15:26:35.0', '2024-05-03 15:26:35.4', '2024-05-03 15:26:35.8',
        '2024-05-03 15:26:36.2', '2024-05-03 15:26:36.6', '2024-05-03 15:26:37.0',
        '2024-05-03 15:26:37.4', '2024-05-03 15:26:37.8', '2024-05-03 15:26:38.2',
        '2024-05-03 15:26:38.6', '2024-05-03 15:26:39.0', '2024-05-03 15:26:39.4',
        '2024-05-03 15:26:39.8', '2024-05-03 15:26:40.2', '2024-05-03 15:26:40.6',
        '2024-05-03 15:26:41.0', '2024-05-03 15:26:41.4', '2024-05-03 15:26:41.8',
        '2024-05-03 15:26:42.2', '2024-05-03 15:26:42.6'
    ]),
    'status': [None]*10 + ['midway'] + [None]*9
}, index=range(235, 255))

期望输出示例 #

df_output = pd.DataFrame({
    'DateTime': pd.to_datetime([
        '2024-05-03 15:26:35.0', '2024-05-03 15:26:35.4', '2024-05-03 15:26:35.8',
        '2024-05-03 15:26:36.2', '2024-05-03 15:26:36.6', '2024-05-03 15:26:37.0',
        '2024-05-03 15:26:37.4', '2024-05-03 15:26:37.8', '2024-05-03 15:26:38.2',
        '2024-05-03 15:26:38.6', '2024-05-03 15:26:39.0', '2024-05-03 15:26:39.4',
        '2024-05-03 15:26:39.8', '2024-05-03 15:26:40.2', '2024-05-03 15:26:40.6',
        '2024-05-03 15:26:41.0', '2024-05-03 15:26:41.4', '2024-05-03 15:26:41.8',
        '2024-05-03 15:26:42.2', '2024-05-03 15:26:42.6'
    ]),
    'status': [None]*4 + ['rest_start', None] + ['rest']*4 + ['midway'] + ['rest']*5 + ['rest_end'] + [None]*3
}, index=range(235, 255))

向量化实现方案 #

步骤1：确保DateTime列为datetime类型 #

先转换DateTime列为datetime格式（如果尚未转换）：

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

步骤2：提取midway时间点并计算区间边界 #

提取所有标记为midway的时间，计算每个midway对应的理论起点（midway - 2.5秒）和理论终点（midway + 2.5秒）：

midway_times = df[df['status'] == 'midway']['DateTime'].reset_index(drop=True)
rest_intervals = pd.DataFrame({
    'midway': midway_times,
    'start_theory': midway_times - pd.Timedelta(seconds=2.5),
    'end_theory': midway_times + pd.Timedelta(seconds=2.5)
})

步骤3：用merge_asof匹配最接近的实际起止行 #

使用Pandas的merge_asof（向量化操作）找到每个理论起止时间对应的最接近的实际行，获取它们的索引：

# 先对原DataFrame按时间排序
df_sorted = df.sort_values('DateTime')

# 匹配rest_start对应的行（找最接近理论起点的记录）
rest_starts = pd.merge_asof(
    rest_intervals[['start_theory']],
    df_sorted.reset_index(),
    left_on='start_theory',
    right_on='DateTime',
    direction='nearest'
)['index'].values

# 匹配rest_end对应的行（找最接近理论终点的记录）
rest_ends = pd.merge_asof(
    rest_intervals[['end_theory']],
    df_sorted.reset_index(),
    left_on='end_theory',
    right_on='DateTime',
    direction='nearest'
)['index'].values

# 整理每个休息周期的起止索引和midway索引
midway_indices = df[df['status'] == 'midway'].index.values
rest_periods = list(zip(rest_starts, midway_indices, rest_ends))

步骤4：向量化赋值状态 #

通过布尔掩码和索引切片完成批量赋值，避免逐行迭代：

# 初始化临时status列，保留原有的midway标记
temp_status = df['status'].copy()

# 遍历每个休息周期（仅循环周期数量，而非DataFrame行）
for start_idx, mid_idx, end_idx in rest_periods:
    # 标记rest_start
    temp_status.loc[start_idx] = 'rest_start'
    # 标记rest_end
    temp_status.loc[end_idx] = 'rest_end'
    # 标记midway和起止之间的rest（排除start、mid、end）
    rest_mask = (df.index > start_idx) & (df.index < mid_idx) | (df.index > mid_idx) & (df.index < end_idx)
    temp_status[rest_mask] = 'rest'

# 将临时列赋值回原DataFrame
df['status'] = temp_status

方案说明 #

• merge_asof是核心的向量化工具，专门用于按时间匹配最近的记录，比逐行查找效率高几个数量级。
• 最后的循环仅针对休息周期的数量（而非DataFrame的每一行），当midway数量不多时，效率几乎等同于完全向量化。
• 该方案支持多个不重叠的休息周期，无需额外修改逻辑。

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