如何对数据进行分层抽样并把抽样和未抽样的数据分开成两个数据集python
时间: 2023-06-04 10:05:37 浏览: 234
可以使用sklearn库中的StratifiedShuffleSplit方法进行分层抽样。具体步骤如下:
1.将数据集按照标签分组,然后每个标签组内进行随机抽样,得到每个标签组的抽样数据集。
2.将每个标签组的抽样数据集合并成一个新的数据集,即为抽样数据集。
3.将原始数据集和抽样数据集通过差集操作,得到未抽样的数据集。
4.分别保存抽样数据集和未抽样的数据集。
具体的代码实现如下:
# 导入所需的库
from sklearn.model_selection import StratifiedShuffleSplit
import pandas as pd
# 读取数据集
data = pd.read_csv('data.csv')
# 定义分层抽样对象
split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
# 进行分层抽样,得到抽样数据集的索引
for train_index, test_index in split.split(data, data["label"]):
strat_train_set = data.loc[train_index]
strat_test_set = data.loc[test_index]
#保存抽样数据集和原始数据集
strat_train_set.to_csv('strat_train_set.csv', index=False)
data.drop(strat_train_set.index, axis=0).to_csv('no_strat_train_set.csv', index=False)
# 输出结果
print('抽样数据集数量:',len(strat_train_set))
print('未抽样数据集数量:',len(data)-len(strat_train_set))
阅读全文