Python Pandas 基于同列前值与列条件的元素填充向量化实现

实现方案 #

以下为全向量化实现逻辑，无循环遍历，时间复杂度O(n)，适配千万级以下数据量的高性能处理：

import pandas as pd
import numpy as np

# 构造示例数据
df = pd.DataFrame({
    'status': [np.nan, 'busy', 'free', 'busy', 'busy', 'busy', 'free', np.nan],
    'date_start': pd.to_datetime([
        '2021-12-06 09:00:00',
        '2021-12-06 09:17:02',
        '2021-12-06 09:18:32',
        '2021-12-06 09:32:45',
        '2021-12-06 09:41:07',
        '2021-12-06 10:08:01',
        '2021-12-06 10:17:00',
        '2021-12-06 10:18:01'
    ])
})

# 核心处理逻辑
# 1. 标记新busy时段的起始行：当前行是busy且前一行不是busy
is_new_busy = (df['status'] == 'busy') & (df['status'].shift() != 'busy')
# 2. 为连续的busy段分配相同的组ID
busy_group = is_new_busy.cumsum()
# 3. 非busy行的组ID设为空，分组时自动忽略
busy_group[df['status'] != 'busy'] = pd.NA
# 4. 同组内取最早的起始时间作为该段所有busy行的date_start_busy
df['date_start_busy'] = df.groupby(busy_group)['date_start'].transform('first')

执行后输出的df完全符合需求预期。

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