Python Pandas按指定条件匹配两个DataFrame并提取目标字段

匹配规则 #

• 关联条件1：两表state字段值完全一致
• 关联条件2：df1.pre_year对应年份的计数值两表相等，例如pre_year=2018时，需满足df1.year_2018 == df2.year_2018
• 输出要求：
  • 完整保留df1所有原始字段
  • 新增匹配到的df2记录ID字段（命名为df2_matched_id，避免和df1原有ID字段冲突）
  • 新增post year计数字段：post year计算规则为pre_year + 2，若计算结果大于2020则赋值为NA，否则取匹配到的df2记录对应年份列的数值
  • 未匹配到df2记录的df1行全部保留，新增字段填充NA

测试数据 #

import pandas as pd
df1 = pd.DataFrame({'ID' : ['100A', '300A'],
                   'state':['GA', 'FL'],
                   'pre_year':[2018, 2020],
                   'year_2018':[10, 5],
                   'year_2019':[9, 2],
                    'year_2020':[7, 6]
                   })

df2 = pd.DataFrame({'ID' : ['500A', '600A', '700A', '800A'],
                   'state':['GA', 'NY', 'FL','GA'],
                   'year_2018':[10, 0,0,10],
                   'year_2019':[0,3,0,4],
                    'year_2020':[0,0,0,1]
                   })

实现代码 #

# 重命名df2字段避免和df1重名
df2_renamed = df2.rename(columns={
    'ID': 'df2_matched_id',
    'year_2018': 'df2_y2018',
    'year_2019': 'df2_y2019',
    'year_2020': 'df2_y2020'
})

# 按state左连接，保留所有df1行
merged = df1.merge(df2_renamed, on='state', how='left')

# 筛选pre_year对应计数值相等的有效匹配
def match_check(row):
    pre_y = row['pre_year']
    return row[f'year_{pre_y}'] == row[f'df2_y{pre_y}']
merged['valid_match'] = merged.apply(match_check, axis=1)
# 无效匹配的df2字段置空
merged.loc[~merged['valid_match'], ['df2_matched_id', 'df2_y2018', 'df2_y2019', 'df2_y2020']] = pd.NA

# 计算post year计数值
def calc_post_val(row):
    if pd.isna(row['df2_matched_id']):
        return pd.NA
    post_y = row['pre_year'] + 2
    if post_y > 2020:
        return pd.NA
    return row[f'df2_y{post_y}']
merged['post_year_count'] = merged.apply(calc_post_val, axis=1)

# 清理临时字段得到最终结果
final_df = merged.drop(columns=['df2_y2018', 'df2_y2019', 'df2_y2020', 'valid_match'])

运行结果 #

ID state  pre_year  year_2018  year_2019  year_2020 df2_matched_id  post_year_count
0  100A    GA      2018         10          9          7           500A              0.0
1  100A    GA      2018         10          9          7           800A              1.0
2  300A    FL      2020          5          2          6           <NA>              NaN

如果同一df1行匹配到多条df2记录，需要保留单条结果的，可以在最后追加final_df = final_df.drop_duplicates(subset=df1.columns, keep='first')即可。

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