百万级DataFrame按Indicator关联的10天周期分组优化方案问询

现有百万行级别的Pandas DataFrame，结构如下：

id  datetime             indicator  other_values ...
1   2020-01-14 00:12:00  0          ...
1   2020-01-17 00:23:00  1          ...
... 
1   2021-02-01 00:00:00  0          ...
2   2020-01-15 00:05:00  0          ...
2   2020-03-10 00:07:00  0          ...
... 
2   2021-05-22 00:00:00  1          ...
... 

数据仅按id和datetime排序，无完整每日数据，但单日可有多条记录。

需求：
对每个indicator==1的条目，收集同id且datetime在其至多10天前的所有行，丢弃其余行；将结果按对应indicator条目分组（如示例中的A、B组），保留other_values等所有字段，用于后续神经网络训练。

示例输入：

id  datetime             indicator  other_values ...
1   2020-01-14 00:12:00  0          ...
1   2020-01-17 00:23:00  1          ...
1   2020-01-17 00:13:00  0          ...
1   2020-01-20 00:05:00  0          ...
1   2020-03-10 00:07:00  0          ...
1   2020-05-19 00:00:00  0          ...
1   2020-05-20 00:00:00  1          ...

期望输出：

id  datetime             group      other_values ...
1   2020-01-14 00:12:00  A          ...
1   2020-01-17 00:23:00  A          ...
1   2020-01-17 00:13:00  A          ...
1   2020-05-19 00:00:00  B          ...
1   2020-05-20 00:00:00  B          ...

要求基于df.groupby('id')、df.sort_values()等Pandas方法实现，避免朴素Python循环以适配百万级数据量。

步骤1：预处理datetime字段 #

确保datetime列为Pandas datetime类型，若未转换先执行：

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

步骤2：按id分组处理，匹配时间窗口 #

利用groupby结合merge_asof实现高效时间匹配，具体逻辑：

1. 提取所有indicator==1的条目作为锚点，为每个锚点生成唯一组标识；
2. 对每个id组，将原始数据与锚点数据按时间进行merge_asof，匹配到最近的、时间差≤10天的锚点；
3. 过滤掉未匹配到锚点的行，保留匹配成功的行并关联组标识。

代码实现：

import pandas as pd

# 1. 提取锚点数据并生成组标识
anchor_df = df[df['indicator'] == 1].copy()
# 按顺序生成A、B、C...组名，也可替换为数字编号提升效率
anchor_df['group'] = [chr(ord('A') + i) for i in range(len(anchor_df))]

# 2. 按id分组执行时间匹配
def process_group(g):
    # 确保组内数据按datetime排序
    g_sorted = g.sort_values('datetime')
    anchors_sorted = anchor_df[anchor_df['id'] == g.name].sort_values('datetime')
    
    # merge_asof：匹配最近的、时间<=当前行且时间差≤10天的锚点
    merged = pd.merge_asof(
        g_sorted,
        anchors_sorted[['id', 'datetime', 'group']],
        on='datetime',
        by='id',
        direction='backward',
        tolerance=pd.Timedelta(days=10)
    )
    # 过滤未匹配到组的行
    return merged.dropna(subset=['group'])

# 应用到所有id组并合并结果
result_df = df.groupby('id', group_keys=False).apply(process_group)

# 3. 调整列顺序（可选，对齐示例输出）
keep_cols = ['id', 'datetime', 'group', 'other_values']
extra_cols = [col for col in df.columns if col not in keep_cols + ['indicator']]
result_df = result_df[keep_cols + extra_cols]

方案优势 #

• 基于Pandas向量化操作和merge_asof的高效时间匹配，避免循环，适配百万级数据；
• merge_asof针对排序后的时间序列做快速匹配，时间复杂度远低于朴素循环；
• 分组处理逻辑清晰，内存占用可控。

补充说明 #

• 若不需要字母组名，直接用数字编号（如anchor_df['group'] = range(1, len(anchor_df)+1)）更高效；
• 确保原始数据和锚点数据在组内已按datetime排序，merge_asof会强制要求排序，未排序会触发报错；
• 同一id下多个锚点时间窗口重叠时，merge_asof会匹配最近的锚点，符合“对应indicator条目分组”的需求逻辑。

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