train_test_split 的 stratify 参数是实现分层抽样的关键，它按类别标签分组抽样以保持各类别在训练集和测试集中比例一致，避免小类样本丢失；需传入1d标签向量，不支持one-hot编码或多输出，且shuffle必须为true。
直接用 train_test_split 的 stratify 参数就能保持训练集和测试集中各类别比例一致，不需要手动计算或重写逻辑。
为什么 stratify 参数是关键
分层抽样的核心是让每个子集的类别分布与原始数据一致。Scikit-learn 的 train_test_split 在启用 stratify=y 后会按类别标签做分组抽样，而不是全局随机打乱——这意味着即使某类样本极少（比如只有 5 个），它也会尽量保证该类在训练/测试中都出现，避免某类完全丢失。
-
stratify必须是与X长度相同的 1D 数组，通常就是你的标签向量y - 如果
y是 pandas Series，需确保其索引与X对齐，否则可能因隐式对齐导致样本错位 - 不支持多输出（如 multi-label 分类中的
y是二维数组），此时需改用StratifiedShuffleSplit
常见错误：抽样后类别比例明显偏差
这通常不是算法问题，而是数据本身或参数配置引发的假象：
- 原始
y中某类样本数少于test_size * 总数，比如总共 12 个样本、test_size=0.3，那么测试集目标是约 4 个样本；若某类只有 3 个，则无法严格按比例分配——train_test_split会优先保最小类的完整性，导致其他类比例微调 - 用了
shuffle=False却期望分层效果：分层依赖随机打乱后的有序分组，关掉 shuffle 就退化为前 N 行进训练、后 M 行进测试，完全失去分层意义 - 传入了 one-hot 编码的
y（如 shape=(n, k)）：stratify只接受 1D 标签，需先用np.argmax(y, axis=1)或y.idxmax(axis=1)转回整数标签
替代方案：StratifiedShuffleSplit 更灵活的场景
当你需要多次生成不同随机划分（比如交叉验证中反复切分），或者想控制每次切分的随机种子而不重复调用 train_test_split，用 StratifiedShuffleSplit 更清晰：
from sklearn.model_selection import StratifiedShuffleSplit sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42) train_idx, test_idx = next(sss.split(X, y)) X_train, X_test = X[train_idx], X[test_idx] y_train, y_test = y[train_idx], y[test_idx]
注意：n_splits=1 是必须指定的，不能省略；next(...) 是获取第一次划分的惯用写法。
真正容易被忽略的是：分层抽样只保证「单次划分」的类别比例接近原分布，不保证多次运行后各次测试集的类别方差稳定——如果你在做严谨的评估，建议配合 StratifiedKFold 做多次折叠，而不是只依赖一次 train_test_split。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
