在python中通过随机分配创建组

Question

我有一个数据集，模型得分分为三类(高、中、低)。下表如下：

| Score   |
| ------- |
| high    |
| high    |
| high    |
| low     |
| low     |
| low     |
| medium  |
| medium  |
| medium  |

我想把这些分数随机分为4组。control、treatment 1、treatment 2、treatment 3。control组应具有20%的观察结果，其余80%必须分为其他3个同等大小的组。然而，我希望每组的分数分布(高、中、低)是相等的。如何使用python解决此问题？

附言：这只是实际表格的一个表示，但它会有更多的观察结果。

Answer 1

您可以尝试groupby.transform:

cats = [ 'control', 'treatment 1', 'treatment 2', 'treatment 3']
probs = [.2, .8/3, .8/3, .8/3]

(df.groupby('Score')['Score']
.transform(lambda x: np.random.choice(cats, size=len(x), p=probs, replace=True)
)