我有一个名为 df 的数据框，格式如下：

    cat_var_1    cat_var_2     num_var_10    Orange       Monkey         341    Banana        Cat           562    Orange        Dog           223    Banana       Monkey          6..

假设数据集中 cat_var_1 的可能值及其比例为 [‘Orange’: 0.6, ‘Banana’: 0.4]，而 cat_var_2 的可能值及其比例为 [‘Monkey’: 0.2, ‘Cat’: 0.7, ‘Dog’: 0.1]。

如何将数据拆分为训练集、测试集和验证集（60:20:20 的比例），同时保持分类变量的比例不变？在实际操作中，这些变量的数量可以是任意数量，而不仅仅是两个。此外，很明显，精确的比例在实践中可能无法完全实现，但我们希望尽可能接近这些比例。

我已经查看了 sklearn 中描述的 StratifiedKFold 方法：如何将数据集拆分为训练集和验证集，同时保持类别之间的比例？，但这种方法仅限于基于一个分类变量进行评估。

另外，如果您能提供您所实现的解决方案的复杂度，我将不胜感激。

回答：

您可以将 df.cat_var_1+ "_" + df.cat_var_2 传递给 StratifiedShuffleSplit.split() 的参数 y：

但这里有一个使用 DataFrame.groupby 的方法：