基于entry/exit标识对Pandas DataFrame特定单元格执行条件计算

我有如下Pandas DataFrame：

amount entry/exit
  1000       exit
  1100      entry
  1140      entry
  1200      entry
  1500       exit
  1400       exit
  1200      entry
  1100      entry
  1000       exit
  1100       exit
  1500      entry

需要实现以下计算逻辑：

• 针对每个连续entry组中的第一个entry，计算其后续第一个exit的amount与该entry的amount的差值（公式：next exit - entry）
• 其余单元格统一标记为ignore
• 共有2000个此类文件需要批量处理，需遵循：仅处理每组连续entry的第一个entry，忽略entry组之间的第一个exit，下一次计算从exit之后的第一个entry开始

最终期望得到的结果如下：

amount entry/exit difference
  1000       exit     ignore
  1100      entry        400
  1140      entry     ignore
  1200      entry     ignore
  1500       exit     ignore
  1400       exit     ignore
  1200      entry       -200
  1100      entry     ignore
  1000       exit     ignore
  1100       exit     ignore
  1500      entry     ignore

步骤1：标记连续的entry/exit分组 #

首先通过shift()和cumsum()生成连续分组的ID，用于区分不同的连续entry或exit段：

import pandas as pd

# 单文件读取示例（批量处理时可循环读取）
df = pd.read_csv('your_file.csv', delim_whitespace=True)

# 生成组ID：当前行与上一行的entry/exit类型不同时，组ID自增
df['group_id'] = (df['entry/exit'] != df['entry/exit'].shift()).cumsum()

步骤2：标记每个entry组的第一个条目 #

筛选出entry组，并标记每组的第一行：

# 标记是否为entry组的首个条目
df['is_first_entry'] = df.apply(
    lambda x: x['entry/exit'] == 'entry' and df[df['group_id'] == x['group_id']].index[0] == x.name,
    axis=1
)

步骤3：匹配首个entry对应的下一个exit并计算差值 #

提取所有exit的位置和金额，用merge_asof匹配每个首个entry之后的第一个exit：

# 提取exit数据，重置索引用于匹配
exit_df = df[df['entry/exit'] == 'exit'][['amount']].reset_index().rename(columns={'index': 'exit_index', 'amount': 'exit_amount'})

# 提取首个entry数据，重置索引用于匹配
first_entry_df = df[df['is_first_entry']][['amount']].reset_index().rename(columns={'index': 'entry_index', 'amount': 'entry_amount'})

# 匹配每个entry之后的第一个exit
matched = pd.merge_asof(
    first_entry_df.sort_values('entry_index'),
    exit_df.sort_values('exit_index'),
    left_on='entry_index',
    right_on='exit_index',
    direction='forward'
)

# 计算差值
matched['difference'] = matched['exit_amount'] - matched['entry_amount']

步骤4：填充结果到原DataFrame #

初始化difference列为ignore，再将计算好的差值填充到对应行：

df['difference'] = 'ignore'

# 将差值写入对应行
for idx, diff in zip(matched['entry_index'], matched['difference']):
    df.loc[idx, 'difference'] = diff

步骤5：批量处理2000个文件 #

用os模块循环读取目录下的所有文件，执行上述逻辑后保存结果：

import os

input_dir = '你的输入目录路径'
output_dir = '你的输出目录路径'

# 创建输出目录（不存在则自动创建）
os.makedirs(output_dir, exist_ok=True)

# 遍历所有csv文件
for filename in os.listdir(input_dir):
    if filename.endswith('.csv'):
        file_path = os.path.join(input_dir, filename)
        # 读取文件
        df = pd.read_csv(file_path, delim_whitespace=True)
        
        # 执行步骤1-4的逻辑
        df['group_id'] = (df['entry/exit'] != df['entry/exit'].shift()).cumsum()
        df['is_first_entry'] = df.apply(
            lambda x: x['entry/exit'] == 'entry' and df[df['group_id'] == x['group_id']].index[0] == x.name,
            axis=1
        )
        exit_df = df[df['entry/exit'] == 'exit'][['amount']].reset_index().rename(columns={'index': 'exit_index', 'amount': 'exit_amount'})
        first_entry_df = df[df['is_first_entry']][['amount']].reset_index().rename(columns={'index': 'entry_index', 'amount': 'entry_amount'})
        matched = pd.merge_asof(
            first_entry_df.sort_values('entry_index'),
            exit_df.sort_values('exit_index'),
            left_on='entry_index',
            right_on='exit_index',
            direction='forward'
        )
        matched['difference'] = matched['exit_amount'] - matched['entry_amount']
        df['difference'] = 'ignore'
        for idx, diff in zip(matched['entry_index'], matched['difference']):
            df.loc[idx, 'difference'] = diff
        
        # 保存处理后的文件
        output_path = os.path.join(output_dir, f'processed_{filename}')
        df.to_csv(output_path, index=False, sep=' ')

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