python pandas：根据百分比随机分配对照组和治疗组

Question

我正在进行一个实验设计，我需要将一个数据帧df按%分成一个对照组和治疗组。在

这是数据帧df：

df.head()

customer_id | Group | many other columns
ABC             1
CDE             1
BHF             2
NID             1
WKL             2
SDI             2

pd.pivot_table(df,index=['Group'],values=["customer_id"],aggfunc=lambda x: len(x.unique()))

Group 1  : 55394
Group 2  : 34889

现在我需要在df中添加一个标记为“Flag”的列。 对于第1组，我想随机分配50%的“控制”和50%的“测试”。 对于第2组，我想随机分配40%的“控制”和60%的“测试”。在

我要寻找的输出：

^{pr2}$

Answer 1

我们可以使用numpy.random.choice()方法：

In [160]: df['Flag'] = \
     ...: df.groupby('Group')['customer_id']\
     ...:   .transform(lambda x: np.random.choice(['Control','Test'], len(x), 
                                                  p=[.5,.5] if x.name==1 else [.4,.6]))
     ...:

In [161]: df
Out[161]:
  customer_id  Group     Flag
0         ABC      1  Control
1         CDE      1     Test
2         BHF      2     Test
3         NID      1  Control
4         WKL      2     Test
5         SDI      2  Control

更新：

^{pr2}$