如何高效标记Pandas DataFrame中符合时序的正负值配对行

问题说明 #

需要标记两类满足条件的行：

• 存在对应负值，且该负值的日期不早于当前正值日期的正值行
• 存在对应正值，且当前负值日期不早于该正值日期的负值行
  核心逻辑是：同一分组（region+productid+vendor）内，基于绝对值匹配正负值，同时满足时序要求。

解决方案1：基于Merge的矢量化匹配 #

逻辑直观，适合中小数据集，完全避免嵌套循环，利用Pandas矢量化操作提升效率：

import pandas as pd
import datetime

# 示例数据
data = [['US', '100', 'Ven1', datetime.datetime(2020, 5, 17), -100], 
        ['US', '100', 'Ven1', datetime.datetime(2020, 5, 19), 100], 
        ['US', '100', 'Ven1', datetime.datetime(2020, 5, 25), -100], 
        ['CA', 'AR-100', '1276238', datetime.datetime(2020, 3, 25), 10], 
        ['UK', 'UKV', 'Daily', datetime.datetime(2022, 1, 12), 2500], 
        ['UK', 'UKV', 'Daily', datetime.datetime(2022, 1, 12), -2500], 
        ['UK', 'UKV', 'Daily', datetime.datetime(2022, 1, 14), 2500]]
  
df = pd.DataFrame(data, columns=['region', 'productid', 'vendor', 'date', 'qty'])

# 步骤1：添加绝对值列用于匹配正负值，保留原始索引方便后续标记
df['abs_qty'] = df['qty'].abs()
df = df.reset_index()

# 步骤2：分离正负数据集
positive = df[df['qty'] > 0]
negative = df[df['qty'] < 0]

# 步骤3：合并找到符合条件的配对：同分组、同绝对值、负值日期≥正值日期
valid_pairs = pd.merge(
    positive,
    negative,
    on=['region', 'productid', 'vendor', 'abs_qty'],
    suffixes=('_pos', '_neg'),
    how='inner'
).query('date_neg >= date_pos')

# 步骤4：收集所有符合条件的行索引
valid_indices = set(valid_pairs['index_pos'].tolist() + valid_pairs['index_neg'].tolist())

# 步骤5：标记结果并还原原始索引
df['is_valid'] = df['index'].isin(valid_indices)
df = df.set_index('index').rename_axis(None)

# 输出结果
print(df[['region', 'productid', 'vendor', 'date', 'qty', 'is_valid']])

输出结果：

region productid   vendor       date   qty  is_valid
0     US       100     Ven1 2020-05-17  -100     False
1     US       100     Ven1 2020-05-19   100      True
2     US       100     Ven1 2020-05-25  -100      True
3     CA    AR-100  1276238 2020-03-25    10     False
4     UK       UKV    Daily 2022-01-12  2500      True
5     UK       UKV    Daily 2022-01-12 -2500      True
6     UK       UKV    Daily 2022-01-14  2500     False

解决方案2：Merge Asof（大数据量更高效） #

如果数据集规模较大，merge_asof是更优选择——它基于排序后的序列做匹配，时间复杂度更低，适合海量数据场景：

import pandas as pd
import datetime

# 示例数据
data = [['US', '100', 'Ven1', datetime.datetime(2020, 5, 17), -100], 
        ['US', '100', 'Ven1', datetime.datetime(2020, 5, 19), 100], 
        ['US', '100', 'Ven1', datetime.datetime(2020, 5, 25), -100], 
        ['CA', 'AR-100', '1276238', datetime.datetime(2020, 3, 25), 10], 
        ['UK', 'UKV', 'Daily', datetime.datetime(2022, 1, 12), 2500], 
        ['UK', 'UKV', 'Daily', datetime.datetime(2022, 1, 12), -2500], 
        ['UK', 'UKV', 'Daily', datetime.datetime(2022, 1, 14), 2500]]
  
df = pd.DataFrame(data, columns=['region', 'productid', 'vendor', 'date', 'qty'])

# 步骤1：按分组键和日期排序（merge_asof要求输入已排序）
df_sorted = df.sort_values(by=['region', 'productid', 'vendor', 'date']).reset_index()
df_sorted['abs_qty'] = df_sorted['qty'].abs()

# 步骤2：分离正负数据集
positive = df_sorted[df_sorted['qty'] > 0]
negative = df_sorted[df_sorted['qty'] < 0]

# 步骤3：匹配正值对应的合格负值（日期≥正值日期）
matches_pos = pd.merge_asof(
    positive,
    negative,
    by=['region', 'productid', 'vendor', 'abs_qty'],
    on='date',
    direction='forward'  # 找当前日期及之后的第一个匹配负值
)

# 步骤4：匹配负值对应的合格正值（日期≤负值日期）
matches_neg = pd.merge_asof(
    negative,
    positive,
    by=['region', 'productid', 'vendor', 'abs_qty'],
    on='date',
    direction='backward'  # 找当前日期及之前的最后一个匹配正值
)

# 步骤5：收集有效索引
valid_pos_indices = matches_pos[~matches_pos['qty_y'].isna()]['index']
valid_neg_indices = matches_neg[~matches_neg['qty_y'].isna()]['index']

# 步骤6：标记结果
df['is_valid'] = df.index.isin(valid_pos_indices) | df.index.isin(valid_neg_indices)

# 输出结果
print(df[['region', 'productid', 'vendor', 'date', 'qty', 'is_valid']])

方案优势 #

• 两种方案均采用Pandas矢量化操作，效率远高于嵌套循环逐行检查
• merge方法逻辑直观，适合中小数据集快速实现
• merge_asof适合大数据量，匹配速度更快，尤其当分组内数据有序时

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