我正在进行实验设计,需要通过预先存在的分组按百分比将数据框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
关于python - python pandas:根据百分比随机分配对照组和治疗组,我们在Stack Overflow上找到一个类似的问题:https://stackoverflow.com/questions/46548404/