Python不使用scikit-learn实现自定义分层抽样方法

不需要依赖scikit-learn，仅用NumPy内置函数即可完成分组按比例打标，逻辑如下：

• 先创建和原数组等长的初始全0输出数组
• 按原数组的不同取值分组，定位每个分组对应的所有元素索引
• 每组内随机打乱索引顺序，按80%的比例取对应位置标记为1，剩余20%位置保持0
• 所有分组处理完成后，输出数组和原数组索引一一对应，满足分层拆分要求

可直接运行的代码 #

import numpy as np

# 原输入数组
Fact = np.array((2,2,2,2,1,2,1,1,2,2,2,1,2,2,2,1,2,2,2,1,2,2,1,1,2,1,2,2,2,2,2,2,1,2,2))
# 初始化等长的全0输出向量
Output = np.zeros_like(Fact, dtype=int)

for group_val in np.unique(Fact):
    # 定位当前分组的所有元素索引
    group_indices = np.where(Fact == group_val)[0]
    # 随机打乱组内索引，实现随机抽样
    np.random.shuffle(group_indices)
    # 计算80%分位的切分位置：1分组共10个样本切到第8位，2分组共25个样本切到第20位，和要求的数量完全匹配
    split_point = int(len(group_indices) * 0.8)
    # 选中的80%样本标记为1
    Output[group_indices[:split_point]] = 1

# 打印结果验证
print(Output)

注意事项 #

• 上述代码每次运行会随机选择组内20%的样本标记为0，如果需要固定拆分结果复现，可以在shuffle操作前添加np.random.seed(任意固定整数)设置随机种子
• 拆分比例计算时自动向下取整，对于样本数乘80%为整数的场景（比如本题100.8=8、250.8=20）可以完全匹配要求的标记数量，如果遇到样本数无法整除的场景，可以根据业务需求调整split_point的取整规则（四舍五入/向上取整）

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