如何在Pandas中抽取包含列最值的正态分布样本

原代码及执行结果 #

原Pandas代码：

d = {"a": [1, 2, 3, 4, 5], "b": [2, 4, 6, 8, 10]}
df = pd.DataFrame(d)

df.sample(n=3, weights='b', random_state=1)

执行后返回：

a   b
3   4   8
4   5   10
0   1   2

用户需求 #

需要抽取必须包含b列最小值和最大值的样本，剩余1个样本按正态分布规则选取，期望结果示例：

a   b
0   1   2
2   3   6
4   5   10

注：移除weights参数后，仍无法保证样本中包含b列的最小值。

实现代码 #

直接用sample没法同时满足“必选最小/最大值”+“剩余样本正态分布采样”，得拆步骤来做：

import pandas as pd
import numpy as np

d = {"a": [1, 2, 3, 4, 5], "b": [2, 4, 6, 8, 10]}
df = pd.DataFrame(d)

# 1. 先把b列最小、最大值对应的行挑出来
min_b_rows = df[df['b'] == df['b'].min()]
max_b_rows = df[df['b'] == df['b'].max()]

# 2. 从剩下的行里按正态分布采样
remaining_df = df[(df['b'] != df['b'].min()) & (df['b'] != df['b'].max())]
# 以剩余行b列的均值为中心生成正态分布权重
mean = remaining_df['b'].mean()
std = remaining_df['b'].std()
# 计算每个行对应的正态概率密度值作为采样权重
norm_weights = np.exp(-((remaining_df['b'] - mean)**2) / (2 * std**2))
# 补采1行（总样本数3，已选2行）
sampled_rest = remaining_df.sample(n=1, weights=norm_weights, random_state=1)

# 3. 合并所有选中的行
final_sample = pd.concat([min_b_rows, sampled_rest, max_b_rows]).reset_index(drop=True)
print(final_sample)

执行结果 #

a   b
0  1   2
1  3   6
2  5  10

说明 #

• 先固定选中最小、最大值行，彻底解决了普通sample无法保证包含这两行的问题
• 剩余行的采样权重基于正态分布生成，越接近b列均值的行被选中的概率越高，符合正态分布采样的要求

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