pandas样本分层抽样（可以自己设置每一个类别抽取多少个样本）

import numpy
import pandas as pd

data = pd.read_csv('data.csv')
numpy.random.seed(seed=2)  # 如果使用相同的seed( )值，则每次生成的随即数都相同，

# 用groupby进行数据的分组,根据班级
# gbr = data.groupby('class')
# 获取分组后gbr的数据
# gbr.groups

# 分层抽样字典定义 组名：数据个数
typicalNDict = {1: 5000, 0: 5000}


# 函数定义
def typicalsamling(group, typicalNDict):
    name = group.name
    n = typicalNDict[name]
    return group.sample(n=n)


#    返回值：抽样后的数据框
result = data.groupby('class').apply(typicalsamling, typicalNDict)
print(result.head())
result.to_csv('smallData.csv', index=False)

我们可以使用pandas中value_counts方法查看每个类别下面具体有多少个样本，检查我们代码是否正确（这个抽样是随机的）

print(df['class'].value_counts())