如何根据熊猫中的给定概率随机选择一行

时间:2020-09-09 18:45:11

标签: python pandas

我有一个这样的数据框:

>>> df = pd.DataFrame([['a',0,0.2],['b',0,0.3],
...                    ['c',0,0.5],
...                    ['a',1,0.4],['b',1,0.3],['c',1,0.3],
...                    ['a',2,0.5],['b',2,0.5]]
...                    ,columns=['place','ID','prob'])
>>> df
place   ID  prob
0   a   0   0.20
1   b   0   0.30
2   c   0   0.50
3   a   1   0.40
4   b   1   0.30
5   c   1   0.30
6   a   2   0.50
7   b   2   0.50

我想使用概率列作为概率质量分布,以随机选择每个ID中的行之一。也就是说,我只想在每个“ ID”中选择一行。输出看起来像这样:

place   ID  prob    choice
0   a   0   0.20    1
1   b   0   0.30    0
2   c   0   0.50    0
3   a   1   0.40    0
4   b   1   0.30    1
5   c   1   0.30    0
6   a   2   0.50    1
7   b   2   0.50    0

实际数据帧将具有数百万行,因此效率越高越好。谢谢!

1 个答案:

答案 0 :(得分:1)

我们可以将您的prob用作DataFrame.sample中的权重。唯一要做的就是在GroupBy中使用它,因为我们想对place中的每个组执行此操作:

sample = df.groupby("ID").apply(lambda x: x.sample(weights=x["prob"]))
choices = sample.reset_index(drop=True, level=0).index
df["choice"] = df.index.isin(choices).astype(int)

  place  ID  prob  choice
0     a   0   0.2       0
1     b   0   0.3       1
2     c   0   0.5       0
3     a   1   0.4       1
4     b   1   0.3       0
5     c   1   0.3       0
6     a   2   0.5       0
7     b   2   0.5       1