python - 根据超过 2 个组的 % 随机分配控制组和治疗组

Question

小 pig 回避我之前的问题 python pandas: assign control vs. treatment groupings randomly based on %

感谢@maxU，我知道如何将随机控制/治疗分组分配给 2 个组；但是如果我有 3 个或更多组呢？

例如:

df.head()

customer_id | Group | many other columns
ABC             1
CDE             3
BHF             2
NID             1
WKL             3
SDI             2
JSK             1
OSM             3
MPA             2
MAD             1

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

Group 1  : 270
Group 2  : 180
Group 3  : 330

当我只有两组时，我有一个很好的答案:

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

但是如果我想这样拆分呢:

• 第 1 组:50% 控制和 50% 测试
• 第 2 组:40% 控制和 60% 测试
• 第 3 组:20% 控制和 80% 测试

@MaxU 的回答很好，但不幸的是分割不准确

d = {1:[.5,.5], 2:[.4,.6], 3:[.2,.8]}

df['Flag'] = df.groupby('Group')['customer_id'] \
             .transform(lambda x: np.random.choice(['Control','Test'], len(x), p=d[x.name]))

当我测试它时，我没有得到精确的拆分。

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

           Control  Treatment
Group 1:    138       132
Group 2:    78        102
Group 3:    79        251

第 1 组应该是 135/135。

Answer 1

In [13]: df
Out[13]:
  customer_id  Group
0         ABC      1
1         CDE      3
2         BHF      2
3         NID      1
4         WKL      3
5         SDI      2
6         JSK      1
7         OSM      3
8         MPA      2
9         MAD      1

In [14]: d = {1:[.5,.5], 2:[.4,.6], 3:[.2,.8]}

In [15]: df['Flag'] = \
    ...: df.groupby('Group')['customer_id'] \
    ...:   .transform(lambda x: np.random.choice(['Control','Test'], len(x), p=d[x.name]))
    ...:

In [16]: df
Out[16]:
  customer_id  Group     Flag
0         ABC      1  Control
1         CDE      3     Test
2         BHF      2     Test
3         NID      1  Control
4         WKL      3  Control
5         SDI      2     Test
6         JSK      1     Test
7         OSM      3     Test
8         MPA      2  Control
9         MAD      1     Test