Pandas中高效处理分组ID：获取各变量首尾有效时间戳

首先给出示例DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'caseid': [1, 1, 1, 2, 2, 2, 3, 3, 3],
    'timestamp': [10, 20, 30, 10, 20, 30, 10, 20, 30],
    'var1': [np.nan, np.nan, np.nan, 10, np.nan, 11, 12, 13, 14],
    'var2': [2., 3., 4., np.nan, 5., 6., np.nan, np.nan, np.nan]
})

需求说明 #

按caseid分组，为每个变量（var1、var2）获取第一个有有效值对应的timestamp和最后一个有有效值对应的timestamp。例如：

• var1在caseid=1时无有效值，结果为None；
• var1在caseid=2时第一个有效timestamp是10，最后一个是30；
• 其余变量同理。
  要求实现时避免循环caseid或逐列调用first_valid_index()，提升运行效率。

高效实现方案 #

利用groupby.apply结合向量化操作，一次性处理所有变量列，避免低效循环：

# 指定需要处理的变量列
var_cols = ['var1', 'var2']

# 分组后批量计算每个变量的首尾有效timestamp
result = df.groupby('caseid').apply(
    lambda g: pd.Series({
        f'{var}_first_ts': g.loc[g[var].notna(), 'timestamp'].iloc[0] if not g[var].isna().all() else None,
        f'{var}_last_ts': g.loc[g[var].notna(), 'timestamp'].iloc[-1] if not g[var].isna().all() else None
        for var in var_cols
    })
).reset_index()

# 将结果中的NaN替换为None（可选）
result = result.replace({np.nan: None})

输出结果 #

caseid var1_first_ts var1_last_ts var2_first_ts var2_last_ts
0       1          None         None            10           30
1       2            10           30            20           30
2       3            10           30          None         None

方案优势 #

• 无需手动循环caseid或逐列处理，充分利用Pandas的分组聚合能力；
• 逻辑清晰，通过notna()过滤无效行后直接取首尾值，避免冗余计算；
• 相比循环实现，在数据量较大时能显著提升运行效率。

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