Pandas 每組保留前N個值，其他設置為0

Question

我有以下 df，改編自這里

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) 到 select 每個組的前 2 個數字和 (2) 將其他值設置為 0。

有什么建議么？ 謝謝。

Answer 1

使用nlargest

df.loc[~df.index.isin(df.groupby('group')['value']\
                        .apply(lambda s: s.nlargest(2))\
                        .index\
                        .get_level_values(1)), 'value'] = 0

---

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

Answer 2

對 boolean 索引使用nlargest每組：

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

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

Output（為清楚起見，作為新列“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

Answer 3

您可以結合使用where()和cumcount()來獲得您期望的結果

df = pd.DataFrame({'group':[1,1,1,2,2,2,3,3,3],
               'value':[1,2,3,1,2,3,4,3,2]})
df.where(df.sort_values('value', ascending = False).groupby('group')['value'].cumcount() + 1 <= 2, 0)

Answer 4

def function1(dd:pd.DataFrame):
    return dd.assign(value=np.where(dd.index.isin(dd.nlargest(2,'value').index),dd.value,0))

df.groupby('group').apply(function1)

出去

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