如何从pandas DataFrame单列筛选相似值匹配两组年龄分布

核心逻辑 #

你要做的是1:1无放回的病例对照年龄最近邻匹配，流程如下：

• 先从原始数据中拆分出病例组（disease=1）和对照组候选池（disease=0）
• 遍历每一个病例样本，从当前可用的对照组候选池中，找到年龄差绝对值最小的样本
• 把匹配到的对照样本从候选池中移除，避免重复匹配
• 合并所有病例和匹配到的对照样本，得到1次匹配后的df1，重复上述流程100次即可得到100份匹配数据集

代码实现 #

首先导入依赖库：

import pandas as pd
import numpy as np

然后完成数据准备和匹配流程：

# 示例数据，可替换为自己的真实数据集
d = {'age': [25,34,45,67,67,44,55,67,26], 'disease': [1,0,1,0,0,0,0,0,0]}
df = pd.DataFrame(data=d)

# 拆分病例组和对照组候选池
case_df = df.query("disease == 1").reset_index(drop=True)
control_pool = df.query("disease == 0").reset_index(drop=True)
match_rounds = 100  # 匹配次数
matched_datasets = []  # 存储100次匹配的结果

def single_round_match(case_data, control_pool):
    available_controls = control_pool.copy()
    matched_controls = []
    for _, case in case_data.iterrows():
        # 计算当前病例和所有可用对照的年龄差
        available_controls['age_diff'] = abs(available_controls['age'] - case['age'])
        # 取年龄差最小的对照，存在多个相同最小差时随机取1个
        closest_control = available_controls.sort_values(by=['age_diff', np.random.rand(len(available_controls))]).iloc[:1].drop(columns='age_diff')
        matched_controls.append(closest_control)
        # 移除已匹配的对照，避免重复使用
        available_controls = available_controls.drop(closest_control.index)
    # 合并病例和对照得到单次匹配结果
    return pd.concat([case_data, pd.concat(matched_controls)], ignore_index=True)

# 执行100次匹配
for _ in range(match_rounds):
    matched_datasets.append(single_round_match(case_df, control_pool))

# 第一次匹配的结果即为df1
df1 = matched_datasets[0]

性能优化建议 #

如果你的数据集规模超过10万行，可以用KDTree优化最近邻查找速度，替换上面的差值计算逻辑即可：

• 导入scipy.spatial.KDTree对对照组年龄建索引
• 单次最近邻查找时间复杂度从O(n)降低到O(logn)，匹配效率提升显著

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