python - Pandas 保留每组前 N 个值，并将其他值设置为 0

Question

我有以下 df，改编自 here

import pandas as pd    
df = pd.DataFrame({'group':[1,1,1,2,2,2,3,3,3],
               'value':[1,2,3,1,2,3,4,3,2]})

我希望得到以下结果:

df1 = pd.DataFrame({'group':[1,1,1,2,2,2,3,3,3],
               'value':[0,2,3,0,2,3,4,3,0]})

逻辑是(1)为每组选择前 2 个数字，(2)将其他值设置为 0。

有什么建议吗？谢谢。

Answer 1

使用nlargest每组 bool 索引:

m = df.groupby('group')['value'].apply(lambda g: g.isin(g.nlargest(2)))

df['value'] = df['value'].where(m, 0)

输出(为了清晰起见，作为新列“value2”)

   group  value  value2
0      1      1       0
1      1      2       2
2      1      3       3
3      2      1       0
4      2      2       2
5      2      3       3
6      3      4       4
7      3      3       3
8      3      2       0