如何修改Pandas DataFrame中两个WITHDRAW之间首个出现的INVEST的值

核心逻辑 #

• 先创建辅助分组列：每次遇到WITHDRAW类型的记录时，分组编号加1，(df['TYPE'] == 'WITHDRAW').cumsum().shift(fill_value=0) 可以完成这个逻辑，加shift是为了把WITHDRAW记录归入前一个分组，确保每两段WITHDRAW之间的所有记录属于同一个分组
• 逐组遍历，找到每个分组里第一条类型为INVEST的记录，将其TYPE字段修改为FIRST INVEST
• 删除辅助分组列即可得到最终结果

完整代码 #

import pandas as pd

# 构造示例数据，实际使用时替换为你自己的数据源读取代码即可
data = {
    'NAME': ['NAME 1', 'NAME 1', 'NAME 1', 'NAME 1', 'NAME 2', 'NAME 2', 'NAME 2', 'NAME 1', 'NAME 1', 'NAME 1', 'NAME 1', 'NAME 1', 'NAME 1'],
    'DATE': ['01/01/2021', '02/01/2021', '03/01/2021', '04/01/2021', '05/01/2021', '06/01/2021', '07/01/2021', '08/01/2021', '11/01/2021', '12/01/2021', '09/01/2021', '10/01/2021', '13/01/2021'],
    'TYPE': ['INVEST', 'INVEST', 'INVEST', 'WITHDRAW', 'INVEST', 'INVEST', 'WITHDRAW', 'INVEST', 'INVEST', 'INVEST', 'INVEST', 'WITHDRAW', 'WITHDRAW'],
    'VALUE': [1000, 500, 501, 1500, 1000, 500, 1500, 1000, 1000, 500, 500, 1500, 1500]
}
df = pd.DataFrame(data)

# 生成分组标记
df['group'] = (df['TYPE'] == 'WITHDRAW').cumsum().shift(fill_value=0)

# 标记每个分组的首个INVEST
def update_first_invest(group):
    first_inv_idx = group[group['TYPE'] == 'INVEST'].first_valid_index()
    if first_inv_idx is not None:
        group.loc[first_inv_idx, 'TYPE'] = 'FIRST INVEST'
    return group

df = df.groupby('group', group_keys=False).apply(update_first_invest)

# 删除辅助列
df.drop('group', axis=1, inplace=True)

# 输出结果
print(df)

运行上述代码得到的输出和你期望的结果完全一致。

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