python pandas : assign control vs. 根据 % 随机处理分组

我正在进行一项实验设计，我需要根据预先存在的分组按百分比将数据框 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% 的“测试”。

我正在寻找的输出:

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

请您参考如下方法:

我们可以使用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 

更新:

In [8]: df 
Out[8]: 
  customer_id  Group 
0         ABC      1 
1         CDE      1 
2         BHF      2 
3         NID      1 
4         WKL      2 
5         SDI      2 
6         XXX      3 
7         XYZ      3 
8         XXX      3 
 
In [9]: d = {1:[.5,.5], 2:[.4,.6], 3:[.2,.8]} 
 
In [10]: df['Flag'] = \ 
    ...: df.groupby('Group')['customer_id'] \ 
    ...:   .transform(lambda x: np.random.choice(['Control','Test'], len(x), p=d[x.name])) 
    ...: 
 
In [11]: df 
Out[11]: 
  customer_id  Group     Flag 
0         ABC      1     Test 
1         CDE      1     Test 
2         BHF      2  Control 
3         NID      1  Control 
4         WKL      2  Control 
5         SDI      2     Test 
6         XXX      3     Test 
7         XYZ      3     Test 
8         XXX      3     Test