怎么在Python中实现数据的随机抽样_利用sample函数进行分层抽样

在数据处理中，随机抽样是个基础操作，但“分层抽样”却常常让人误入歧途。很多人第一反应是去找一个叫“分层抽样”的函数，结果发现Python内置的random.sample()并不支持。这其实是个典型的认知偏差：我们混淆了“功能需求”和“实现方法”。

sample函数本身不支持分层抽样

必须明确一点：random.sample()是一个纯粹的简单随机抽样工具。它处理的是一个扁平的序列，比如一个列表或元组，然后从中无差别地抽取指定数量的元素。它不具备识别数据内部结构（比如类别、分组）的能力。如果你把不同类别的数据混在一起丢给它，它只会一视同仁地随机抓取，最终结果很可能导致某些类别被严重过采样或欠采样，甚至完全漏掉。这并非函数本身的缺陷，而是方法选错了。

分层抽样的正确逻辑链条是清晰的：

• 首先，必须依据某个关键字段（例如“部门”、“产品类别”、“用户标签”）将数据分成不同的组。
• 然后，在每个独立的组内，分别进行随机抽样，可以控制固定数量或按比例抽取。
• 最后，将各个组抽出的样本合并起来，形成最终的分层样本集。

pandas DataFrame中用groupby + sample实现分层抽样

对于结构化数据，最直观高效的方法是使用pandas。假设你有一个DataFramedf，其中“group”列是你的分层依据。

import pandas as pd
# 按 group 列分层，每组抽 5 行
stratified_sample = df.groupby("group", group_keys=False).apply(lambda x: x.sample(5))
# 或按比例抽（比如每组抽 20%）
stratified_sample = df.groupby("group", group_keys=False).apply(lambda x: x.sample(frac=0.2))

这里有三个关键细节需要注意：

• group_keys=False这个参数至关重要。如果设为True（默认），结果的多级索引里会包含分组键，给后续处理平添麻烦。
• frac（比例）和n（固定数量）不能同时指定。使用frac时要注意，pandas会将其转换为整数（向下取整）。如果某组数据量很少，比如只有3行，设置frac=0.1意味着要抽0.3行，取整后为0，程序就会报错ValueError: Cannot take a larger sample than population when 'replace=False'。
• 默认情况下，sample是不放回抽样。如果某组的行数少于你指定的抽样数量，同样会触发上述错误。虽然可以设置replace=True启用有放回抽样，但这通常不符合分层抽样的初衷。稳妥的做法是先用df["group"].value_counts()检查一下各层的基数。

用scikit-learn的train_test_split做带分层的划分

如果你的目标非常明确，就是为了划分机器学习中的训练集和测试集，并且希望保持类别比例一致，那么sklearn的train_test_split函数提供了更直接的方案。

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
    X, y, 
    test_size=0.3, 
    stratify=y,  # 关键参数：根据y进行分层划分
    random_state=42
)

这个方法底层自动完成了分层逻辑，能确保训练集和测试集中各类别的比例与原始数据集基本一致。但它主要服务于建模流程，输出的是特征矩阵和标签数组，不保留原始的DataFrame索引和结构。因此，它不适合需要自定义每层样本数量或需要保留完整行信息的场景。

使用时需留意：

• stratify参数需要传入一个一维数组（如标签列y），其长度必须与特征数据X的行数一致。
• 如果某个类别的样本数极少，而划分比例又可能导致该类别在某个子集中数量为0，函数会报错ValueError: The least populated class in y has only 1 member。这时需要考虑合并稀有类别或调整采样策略。

手动分组+sample时容易忽略的索引问题

使用groupby().apply(sample)得到的结果，默认会保留原始数据的索引。这在需要回溯原始记录时是优点，但如果你将结果作为一个独立的新数据集进行后续操作，可能会遇到重复索引或排序混乱的问题。

• 一个简单的处理方法是抽样后立即重置索引：stratified_sample.reset_index(drop=True)。
• 但如果原始索引本身具有业务含义（如订单ID、时间戳），则不应重置，而应显式地将索引作为一列数据保存下来。
• 还有一个更隐蔽的情况：在某些pandas版本中，如果某个分组为空，或者sample操作意外返回了一个空的DataFrame，整个apply过程可能会失败或返回异常结果。建议在分组前检查并过滤掉空组，或者在apply函数内增加容错判断。

说到底，分层抽样的技术实现并不复杂，真正的难点在于前期的业务逻辑梳理：依据哪一列来分层？每层是抽取固定数量还是按比例抽取？是否允许某些层因为数据量太少而抽样失败？把这些想清楚了，代码无非就是用groupby还是train_test_split的选择题罢了。核心永远是先理解数据，再操作数据。