Python中如何基于第三变量分箱提取两变量配对数据

将落在第三个变量Z分箱区间内的X、Y两个变量值做配对分组，提取Z特定分箱范围内的XY配对数据，效果参考如下示例：

• 整体配对分组效果：
• 25-50分箱对应的XY配对效果：

已通过如下代码完成Z变量排序，同步对齐对应X、Y变量的顺序，该步骤运行无异常：

Z, Y, X = zip(*sorted(zip(Z,Y,X)))

• 将排序后的Z变量按25为间隔分箱，分箱区间设置为bins=[0,25,50,75,100,125,150,200, 所有大于200的数值]
• 支持灵活调用Z变量任意特定分箱内对应的X、Y配对数据
• 测试数据集获取地址：测试数据

采用pandas实现分箱与分组映射，逻辑简单易维护，代码如下：

import pandas as pd
import numpy as np

# 已排序对齐的Z、Y、X可直接复用之前的运行结果
# 构造结构化数据表存储对齐后的三个变量
df = pd.DataFrame({
    "X": X,
    "Y": Y,
    "Z": Z
})

# 定义分箱边界，np.inf代表正无穷，覆盖所有大于200的数值
bins = [0, 25, 50, 75, 100, 125, 150, 200, np.inf]
# 定义和分箱一一对应的标签，用于后续调用
bin_labels = ["0-25", "25-50", "50-75", "75-100", "100-125", "125-150", "150-200", ">200"]

# 为每条数据匹配所属分箱
df["bin_tag"] = pd.cut(
    df["Z"],
    bins=bins,
    labels=bin_labels,
    include_lowest=True  # 确保Z=0的样本被归入0-25区间
)

# 生成分箱到XY配对数据的映射字典
bin_dataset = {}
for tag in bin_labels:
    subset = df[df["bin_tag"] == tag]
    bin_dataset[tag] = (subset["X"].to_numpy(), subset["Y"].to_numpy())

调用方式 #

需要获取指定分箱的XY配对数据时，直接通过标签从映射字典中取值即可，返回的X、Y数组索引一一对应，即为配对关系。例如获取25-50分箱的配对数据：

x_25to50, y_25to50 = bin_dataset["25-50"]

该方案不会遗漏边界值和大于200的样本，后续如果需要调整分箱间隔，仅需修改bins和bin_labels的配置即可，不需要调整后续分组逻辑。

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