Python中如何洗牌CSV并均匀分布Prediction=1的少数类数据？

你的问题很典型——当数据集严重不平衡（这里Prediction=1的样本仅占1%）时，普通的随机洗牌很容易让少数类样本集中在数据集的某一段，后续划分训练/测试集时可能出现某一集合里几乎没有少数类样本的情况，严重影响模型效果。下面给你两种可行的解决方案：

方案一：纯Python标准库实现（无需额外依赖） #

这种方法通过手动拆分多数类样本，将少数类样本均匀插入到数据集中，保证分布的均匀性：

import csv
import random

# 读取原始数据集
with open("dataset.csv", "r") as f:
    reader = csv.reader(f)
    header = next(reader)
    rows = list(reader)

# 分离少数类（Prediction=1）和多数类（Prediction=0）样本
positive_rows = [row for row in rows if row[2] == "1"]
negative_rows = [row for row in rows if row[2] == "0"]

# 计算每个少数类样本对应的多数类样本数量（按1%的比例）
samples_per_positive = len(negative_rows) // len(positive_rows)
# 将多数类样本切成和少数类数量一致的块
negative_chunks = [negative_rows[i*samples_per_positive : (i+1)*samples_per_positive] 
                   for i in range(len(positive_rows))]
# 处理切割后剩余的多数类样本
remaining_negatives = negative_rows[len(positive_rows)*samples_per_positive :]

# 构建均匀分布的数据集：每个多数类块后插入一个少数类样本
shuffled_data = []
for chunk, pos_row in zip(negative_chunks, positive_rows):
    shuffled_data.extend(chunk)
    shuffled_data.append(pos_row)
# 添加剩余的多数类样本
shuffled_data.extend(remaining_negatives)

# 可选：最后再做一次随机洗牌，避免过于规整的排列
random.shuffle(shuffled_data)

# 保存洗牌后的数据集
with open("dataset_shuffled.csv", "w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(header)
    writer.writerows(shuffled_data)

方案二：使用Scikit-learn的分层洗牌（更简洁） #

如果你可以使用机器学习库scikit-learn，它提供了现成的分层洗牌工具，能自动保持类别比例的均匀分布：

import csv
from sklearn.model_selection import StratifiedShuffleSplit

# 读取原始数据集
with open("dataset.csv", "r") as f:
    reader = csv.reader(f)
    header = next(reader)
    rows = list(reader)

# 提取标签列，用于分层洗牌的依据
labels = [row[2] for row in rows]

# 初始化分层洗牌拆分器：n_splits=1表示只生成一次拆分，test_size=0表示取全部数据
stratified_shuffler = StratifiedShuffleSplit(n_splits=1, test_size=0.0, random_state=42)
# 获取洗牌后的样本索引
for train_indices, _ in stratified_shuffler.split(rows, labels):
    shuffled_rows = [rows[idx] for idx in train_indices]

# 保存结果
with open("dataset_shuffled.csv", "w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(header)
    writer.writerows(shuffled_rows)

为什么这两种方法有效？ #

普通的random.shuffle()是完全随机的，对于极度不平衡的数据集，很容易出现少数类样本扎堆的情况。而上面的方法：

• 方案一通过手动穿插的方式，强制让少数类样本均匀分布在多数类样本中；
• 方案二的StratifiedShuffleSplit会在洗牌过程中保持每个类别在数据集中的分布比例，确保少数类样本不会集中出现。

后续你划分训练/测试集时，也建议使用分层划分（比如sklearn.model_selection.train_test_split的stratify参数），这样能保证训练集和测试集都保持和原数据集一致的类别比例。

内容的提问来源于stack exchange，提问作者Vincent Teyssier