我有一个这样的数据框:
>>> 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
实际数据帧将具有数百万行,因此效率越高越好。谢谢!
答案 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