分层随机样本集划分python
时间: 2024-02-28 09:51:29 浏览: 31
分层随机样本集划分是一种常用的数据集划分方法,它可以保证每个类别在训练集和测试集中的比例相同。在Python中可以使用scikit-learn库中的StratifiedShuffleSplit类来实现分层随机样本集划分。
下面是一个示例代码,演示如何使用StratifiedShuffleSplit进行数据集划分:
```python
from sklearn.model_selection import StratifiedShuffleSplit
# 假设有一个特征矩阵X和对应的目标变量y
X = ...
y = ...
# 创建StratifiedShuffleSplit对象,指定划分比例和随机种子
split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
# 使用split对象进行数据集划分
for train_index, test_index in split.split(X,):
X_train, X_test = X[train_index], X[test_index]
y_train, y_test = y[train_index], y[test_index]
`
在上述代码中,我们首先导入了StratifiedShuffleSplit类。然后,我们创建了一个StratifiedShuffleSplit对象,通过指定参数n_splits为1表示只进行一次划分,test_size为测试集的比例(这里设置为0.2表示测试集占总样本的20%),random_state为随机种子用于重现随机结果。
接下来,我们使用split对象的split方法进行数据集划分。该方法接受特征矩阵X和目标变量y作为输入,并返回训练集和测试集的索引。我们可以通过这些索引从原始数据中获取划分后的训练集和测试集。
需要注意的是,上述代码中的X和y是示例数据,实际使用时需要替换为自己的数据。