如何用Pandas为同Segment下A组数据随机采样B组购买价格？

问题分析 #

你遇到的NaN问题，大概率是因为没有正确按segment分组提取对应B组的价格列表，采样时无法匹配到对应分组的有效数据导致的。下面是直接可用的解决方案：

完整代码实现 #

import pandas as pd
import numpy as np

# 1. 构造原始数据
data = {
    'segment': ['High', 'High', 'High', 'High', 'Low', 'Low', 'Low', 'Low'],
    'purchase price': [100, 105, 103, 104, 10, 9, 50, 55],
    'group': ['A', 'A', 'B', 'B', 'A', 'B', 'B', 'B']
}
df = pd.DataFrame(data)

# 2. 生成每个segment对应的B组价格映射字典
b_group_map = df[df['group'] == 'B'].groupby('segment')['purchase price'].apply(list).to_dict()

# 3. 定义采样逻辑：仅对A组行按对应segment采样B组价格
def get_b_sample(row):
    if row['group'] == 'A':
        # 从当前segment的B组价格列表中随机选一个
        return np.random.choice(b_group_map[row['segment']])
    return np.nan  # 非A组行返回NaN，可根据需求调整

# 4. 添加新列并筛选出A组结果
df['sample_group_B'] = df.apply(get_b_sample, axis=1)
result = df[df['group'] == 'A'].reset_index(drop=True)

print(result)

代码说明 #

• 先通过分组提取每个segment下所有B组的购买价格，存成字典映射，避免重复查询
• 使用np.random.choice()从对应分组的价格列表中随机选取值，确保采样的是当前segment下的B组数据
• 最后筛选出A组行，得到你需要的结果

注意事项 #

• 如果需要固定采样结果（方便调试），可以在代码开头添加np.random.seed(42)设置随机种子
• 如果某个segment没有对应的B组数据，采样时会报错，可提前添加判断逻辑处理这种情况

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