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

Question

我有一个如下所示的 Pandas 数据框:

| 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 的插入 = 1，依此类推。
我查过 this answer ，但它只对一个变量进行分层，它不会对多个变量起作用。
R 有 this function确实如此，但使用 R 不是一种选择。
顺便说一下，我使用的是 Python 3。

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"]])