python - Pandas 基于多列的分层抽样

Question

我有一个看起来像这样的熊猫数据框：

| Cliid | Segment | Insert |
|-------|---------|--------|
| 001   | A       | 0      |
| 002   | A       | 0      |
| 003   | C       | 0      |
| 004   | B       | 1      |
| 005   | A       | 0      |
| 006   | B       | 0      |

我想将其分成 2 组，使每个组在 [Segment, Insert] 中具有相同的每个变量的组成。例如，每个组将有 1/2 的观测值属于分段 A，1/6 的 Insert = 1，依此类推。

我已经检查了这个答案，但它只对一个变量进行分层，它不会对多个变量起作用。

R 有这个功能可以做到这一点，但使用 R 不是一种选择。

顺便说一句，我使用的是 Python 3。

标签： pythonpandas

Answer 1

您可以使用 sklearn 的train_test_split函数，包括stratify可用于确定要分层的列的参数。

例如：

from sklearn.model_selection import train_test_split

df_train, df_test = train_test_split(df1, test_size=0.2, stratify=df[["Segment", "Insert"]])