分层抽样 sklearn

分层抽样(Stratified Sampling)是一种保持样本分布比例的抽样方法。在sklearn中，可以使用StratifiedShuffleSplit类来进行分层抽样。该类可以将数据集划分为训练集和测试集，以便进行模型训练和评估。

使用StratifiedShuffleSplit类进行分层抽样的步骤如下：
1. 首先，导入相关的库和模块，例如导入StratifiedShuffleSplit类和pandas库。
2. 读取数据并进行相应的数据预处理，例如对数据进行特征工程和划分类别。
3. 创建一个StratifiedShuffleSplit对象，并设置参数，如测试集大小和随机种子。
4. 使用split方法对数据集进行分层抽样，该方法返回训练集和测试集的索引。
5. 根据返回的索引，从原始数据集中提取训练集和测试集。

下面是一个示例代码，展示了如何使用StratifiedShuffleSplit类进行分层抽样：

from sklearn.model_selection import StratifiedShuffleSplit
import pandas as pd

# 读取数据并进行数据预处理
df = pd.read_csv('housing.csv')
df['income_cat'] = pd.cut(df['median_income'], bins=[0, 1.5, 3.0, 4.5, 6., np.inf], labels=[1, 2, 3, 4, 5])

# 创建StratifiedShuffleSplit对象
split = StratifiedShuffleSplit(test_size=0.2, random_state=42)

# 使用split方法进行分层抽样
for train_index, test_index in split.split(df, df['income_cat']):
    start_train_set = df.loc[train_index] # 提取训练集
    start_test_set = df.loc[test_index] # 提取测试集