python分层抽样代码
时间: 2023-11-12 16:58:12 浏览: 39
以下是 Python 中的分层抽样代码示例:
```python
import random
def stratified_sampling(data, labels, ratio):
# 创建字典,用于存储每个类别的数据
class_dict = {}
for i in range(len(data)):
if labels[i] not in class_dict:
class_dict[labels[i]] = []
class_dict[labels[i]].append(data[i])
# 计算每个类别需要抽取的样本数
sample_nums = {}
for key in class_dict:
sample_nums[key] = int(len(class_dict[key]) * ratio)
# 抽取样本
samples = []
for key in class_dict:
if sample_nums[key] >= len(class_dict[key]):
samples.extend(class_dict[key])
else:
samples.extend(random.sample(class_dict[key], sample_nums[key]))
return samples
```
其中,`data` 是原始数据,`labels` 是对应的标签,`ratio` 是抽样比例。函数返回抽样后的数据。