如何在Python中嵌套随机选取DataFrame行并赋值（转SAS逻辑）

核心需求 #

完全基于随机数的嵌套逻辑实现分层赋值，不依赖已生成的目标列作为筛选条件，支持多层级嵌套（如某子集内再选子子集赋值），示例场景：

• 第一层：50%行col1设为'A'，剩余50%设为'B'
• 第二层：在第一层分配为'A'的行中，40%设col2为'aa'，60%设'aaa'
• 第三层：在第二层分配为'aa'的行中，20%设col3为'x'，剩余设'y'

SAS参考逻辑 #

SAS中通过嵌套do块结合随机数区间判断实现，完全不依赖生成后的列：

data df;
    set original_df;
    /* 第一层随机分配col1 */
    ran1 = ranuni(0);
    if ran1 <= 0.5 then col1 = 'A';
    else col1 = 'B';
    /* 第二层：仅在ran1<=0.5的行（即col1='A'的行）分配col2 */
    if ran1 <= 0.5 then do;
        ran2 = ranuni(0);
        if ran2 <= 0.4 then col2 = 'aa';
        else col2 = 'aaa';
    end;
    /* 第三层：在ran1<=0.5且ran2<=0.4的行分配col3 */
    if ran1 <=0.5 and ran2 <=0.4 then do;
        ran3 = ranuni(0);
        if ran3 <=0.2 then col3 = 'x';
        else col3 = 'y';
    end;
run;

Python实现（Pandas + NumPy） #

复刻SAS的嵌套随机数逻辑，所有判断基于原始随机数区间，不依赖生成后的目标列：

import pandas as pd
import numpy as np

# 生成测试DataFrame（替换为你的原始数据）
np.random.seed(42)  # 固定种子确保结果可复现
df = pd.DataFrame({'original_data': range(1000)})

# 第一层：生成随机数ran1，分配col1
df['ran1'] = np.random.rand(len(df))
df['col1'] = np.where(df['ran1'] <= 0.5, 'A', 'B')

# 第二层：仅在ran1<=0.5的行，生成ran2并分配col2
df['ran2'] = np.random.rand(len(df))
df['col2'] = np.where(
    df['ran1'] <= 0.5,  # 用ran1判断，而非已生成的col1
    np.where(df['ran2'] <= 0.4, 'aa', 'aaa'),
    np.nan  # 非目标行设为缺失，可替换为其他默认值
)

# 第三层：在ran1<=0.5且ran2<=0.4的行，生成ran3并分配col3
df['ran3'] = np.random.rand(len(df))
df['col3'] = np.where(
    (df['ran1'] <= 0.5) & (df['ran2'] <= 0.4),
    np.where(df['ran3'] <= 0.2, 'x', 'y'),
    np.nan
)

# 可选：删除临时随机数列，保留目标列
df = df.drop(['ran1', 'ran2', 'ran3'], axis=1)

关键说明 #

1. 纯随机数驱动：所有分层判断都基于生成的ran1/ran2/ran3等随机数的区间，完全不依赖已生成的col1/col2，和SAS逻辑完全对齐
2. 无限层级扩展：新增层级只需添加对应的随机数列和多条件np.where判断即可
3. 可验证性：可通过随机数区间直接统计比例，确保符合预期：

# 验证col1的分配比例
print(df['col1'].value_counts(normalize=True))
# 验证col1='A'时col2的分配比例（用ran1判断而非col1）
print(df.loc[df['ran1'] <= 0.5, 'col2'].value_counts(normalize=True))
# 验证col3的分配比例
print(df.loc[(df['ran1'] <= 0.5) & (df['ran2'] <= 0.4), 'col3'].value_counts(normalize=True))

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