如何基于过滤条件对DataFrame采样实现向量化优化？

问题背景 #

我有两个DataFrame：

• df1包含三个变量：离散变量A、连续变量B和C，行数至少10万
• df2包含相同的A、B、C变量，额外多一个变量D

需要给df1添加D列，规则是：对df1每一行，在df2中筛选出A值相等、B在当前行B的±5范围内、C在当前行C的±10范围内的行，再从这些行的D中随机取一个值赋值给df1对应行。

当前用循环遍历df1的方法功能可行，但速度太慢，想找向量化优化方案，考虑过对三个变量分箱预组装过滤后的DataFrame，希望得到可行性建议及其他优化思路。

当前实现代码（含一处逻辑错误） #

import pandas as pd
import numpy as np
import random

# df1
varA1 = [random.randint(1,3) for i in range(1000)]
varB1 = [random.uniform(0, 50) for i in range(1000)]
varC1 = [random.uniform(0, 100) for i in range(1000)]
df1 = pd.DataFrame({'A':varA1,'B':varB1,'C':varC1})

# df2 - reference, look-up.
varA2 = [random.randint(1,3) for i in range(1000)]
varB2 = [random.uniform(0, 50) for i in range(1000)]
varC2 = [random.uniform(0, 100) for i in range(1000)]
varD2 = [random.uniform(0, 1000) for i in range(1000)]
df2 = pd.DataFrame({'A':varA2,'B':varB2,'C':varC2,'D':varD2})

# Assign a value for "D" in Df1, based on values in Df2, considering similar values of other parameters.
df1['D'] = np.nan
for r in df1.iterrows():
    temp = df2.loc[(df2['A']==r[1]['A'])& \
               (df2['B']>=r[1]['B']-5)& \
               (df2['B']<=r[1]['B']+5)& \
               (df2['C'] >= r[1]['C']-10) & \
               # 此处逻辑错误：应为<= r[1]['C']+10
               (df2['C'] >= r[1]['C']+10),:]
    df1.loc[r[0],'D'] = temp['D'].sample().values

注：原代码中C的筛选条件存在错误，第二个比较应为<= r[1]['C']+10，否则会筛选不到符合条件的样本。

优化方案 #

方案一：分箱预分组（你的思路落地） #

你的分箱预分组思路完全可行，且能大幅提升处理速度。核心是提前将df2按离散变量A+连续变量的分箱键分组，再让df1的每行直接匹配对应分组采样，避免重复过滤。

代码实现：

import pandas as pd
import numpy as np

# 预处理df2：生成分箱键并分组保存D列表
# B的分箱：覆盖0-50，步长10，左闭右开
df2['B_bin'] = pd.cut(df2['B'], bins=np.arange(-5, 55, 10), right=False)
# C的分箱：覆盖0-100，步长20，左闭右开
df2['C_bin'] = pd.cut(df2['C'], bins=np.arange(-10, 110, 20), right=False)

# 按A+B_bin+C_bin分组，将每个分组的D转为列表存为字典
grouped_d = df2.groupby(['A', 'B_bin', 'C_bin'])['D'].agg(list).to_dict()

# 处理df1：生成对应分箱键
df1['B_bin'] = pd.cut(df1['B'], bins=np.arange(-5, 55, 10), right=False)
df1['C_bin'] = pd.cut(df1['C'], bins=np.arange(-10, 110, 20), right=False)

# 定义采样函数：匹配分组后随机取一个D值
def sample_d(row):
    key = (row['A'], row['B_bin'], row['C_bin'])
    d_list = grouped_d.get(key, [])
    return np.random.choice(d_list) if d_list else np.nan

# 批量应用采样（比循环快10~100倍）
df1['D'] = df1.apply(sample_d, axis=1)

# 清理临时分箱列
df1.drop(['B_bin', 'C_bin'], axis=1, inplace=True)
df2.drop(['B_bin', 'C_bin'], axis=1, inplace=True)

优缺点：

• 优点：速度最快，内存占用低，适合大数据量场景。
• 缺点：分箱边界需要根据数据范围调整，若样本刚好落在分箱边界，可能出现漏匹配（可通过缩小分箱步长或扩大分组范围解决）。

方案二：KDTree近邻查询（高精度匹配） #

如果对匹配精度要求高，不想用分箱近似，可以用KDTree快速查找符合范围的样本，避免循环过滤。

代码实现：

import pandas as pd
import numpy as np
from scipy.spatial import KDTree

# 按A分组，为每个A的df2数据构建KDTree和D数组
tree_dict = {}
for a_val, group in df2.groupby('A'):
    coords = group[['B', 'C']].values
    tree = KDTree(coords)
    tree_dict[a_val] = (tree, group['D'].values)

# 定义采样函数：用KDTree查询符合范围的样本并采样
def sample_d_kdtree(row):
    a_val = row['A']
    if a_val not in tree_dict:
        return np.nan
    tree, d_vals = tree_dict[a_val]
    # 定义查询矩形范围：[B_min, C_min] 到 [B_max, C_max]
    rect = [[row['B']-5, row['C']-10], [row['B']+5, row['C']+10]]
    # 获取所有符合范围的样本索引
    indices = tree.query_ball_rect(rect)
    return np.random.choice(d_vals[indices]) if indices else np.nan

# 批量应用采样
df1['D'] = df1.apply(sample_d_kdtree, axis=1)

优缺点：

• 优点：无需分箱，匹配精度高，适合对连续变量匹配要求严格的场景。
• 缺点：速度比分箱方案稍慢，但仍远快于循环。

方案选择建议 #

• 若追求极致速度，优先选分箱预分组方案，调整分箱步长即可平衡精度和速度。
• 若要求高精度匹配，选KDTree方案。
• 原循环方案仅适合小体量df1，大数据量下务必替换为上述向量化方案。

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