Python按条件生成date_open和date_close列的实现方法

直接用Pandas的分组功能就能高效解决，比循环写法简洁且性能更好，具体步骤如下：

1. 导入依赖并加载数据
   先把你的数据集转换成Pandas DataFrame：

import pandas as pd

# 构造你的数据集
data = {
    'group': ['A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'C', 'C', 'C', 'C', 'C'],
    'date': ['Jan-22', 'Feb-22', 'Mar-22', 'Apr-22', 'May-22', 'Jun-22',
             'Mar-23', 'Apr-23', 'May-23', 'Jun-23',
             'Sep-21', 'Oct-21', 'Nov-21', 'Dec-21', 'Jan-22'],
    'open': [2,2,2,1,1,0,3,3,3,3,1,1,1,1,0],
    'close': [0,0,0,1,1,2,0,0,0,1,0,0,0,0,1],
    'exit': [0,0,0,0,0,1,0,0,0,0,0,0,0,0,1]
}
df = pd.DataFrame(data)

2. 计算date_open列
   对每个group取第一条记录的date，并把这个值填充到该组的所有行：

df['date_open'] = df.groupby('group')['date'].transform('first')

3. 计算date_close列
   先判断每个组是否存在exit=1的记录，存在则取该组最后一条记录的date，否则设为pd.NA（或0，按需替换）：

# 先分组计算每个组的date_close值
close_dates = df.groupby('group').apply(
    lambda x: x['date'].iloc[-1] if x['exit'].any() else pd.NA
).reset_index(name='date_close')
# 合并回原DataFrame
df = df.merge(close_dates, on='group', how='left')

4. 最终结果
   处理后的数据集如下：

如果需要把<NA>替换成0，只需要在合并后加一行：

df['date_close'] = df['date_close'].fillna(0)

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