优化Pandas DataFrame转换：实现买卖方ID关联的高效方案

原始数据集 #

import pandas as pd

df = pd.DataFrame({
    'SequenceNumber ': [0, 0, 0, 0, 1, 1, 1], 
    'ID': [0, 1, 2, 3, 0, 1, 2], 
    'CountNumber': [3, 1, 1, 1, 1, 1, 2],
    'Side': ['Sell', 'Buy', 'Buy', 'Buy', 'Sell', 'Sell', 'Buy'],
    'featureA': [4, 12, 1, 3, 5, 7, 5],
    'featureB': [2, 45, 4, 36, 11, 12, 35]
})

对应表格：

需求说明 #

按SequenceNumber分组处理：

• 仅保留组内CountNumber == 1的行
• 根据Side生成From/To列：Buy对应To列，Sell对应From列
• 用同组内CountNumber > 1的行的ID填充From/To的空值（每组仅存在1个此类行）
• 保留featureA、featureB特征列

预期输出 #

高效实现方案（无循环，矢量化操作） #

利用Pandas内置的分组、合并和矢量化操作替代循环，大幅提升大数据处理效率：

# 1. 提取每组中CountNumber>1的行，作为填充基准数据
base_df = df[df['CountNumber'] > 1].rename(columns={'ID': 'BaseID', 'Side': 'BaseSide'})
# 生成对应填充列：Sell对应From填充值，Buy对应To填充值
base_df['FillFrom'] = base_df.apply(lambda x: x['BaseID'] if x['BaseSide'] == 'Sell' else None, axis=1)
base_df['FillTo'] = base_df.apply(lambda x: x['BaseID'] if x['BaseSide'] == 'Buy' else None, axis=1)
base_df = base_df[['SequenceNumber ', 'FillFrom', 'FillTo']]

# 2. 提取待处理的CountNumber==1的行，生成初始From/To列
detail_df = df[df['CountNumber'] == 1].copy()
detail_df['From'] = detail_df.apply(lambda x: x['ID'] if x['Side'] == 'Sell' else None, axis=1)
detail_df['To'] = detail_df.apply(lambda x: x['ID'] if x['Side'] == 'Buy' else None, axis=1)

# 3. 合并基准数据与明细数据，填充空值
result_df = detail_df.merge(base_df, on='SequenceNumber ', how='left')
result_df['From'] = result_df['From'].fillna(result_df['FillFrom'])
result_df['To'] = result_df['To'].fillna(result_df['FillTo'])

# 4. 整理最终列并重置索引
result_df = result_df[['SequenceNumber ', 'From', 'To', 'featureA', 'featureB']].reset_index(drop=True)

print(result_df)

方案说明 #

• 拆分基准/明细数据：将每组唯一的大额交易行单独提取，作为统一填充源
• 矢量化赋值：用apply结合条件判断生成初始列，避免逐行循环
• 合并填充：通过merge匹配同组填充值，fillna完成空值填充
• 全程使用Pandas内置矢量化操作，处理百万级数据时性能远优于循环实现

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