如何在 pandas 中的 `groupby` 之后按順序添加一列？

Question

我的數據就像我想計算NDCG指標

  pid query title  label  predict  prob
0   1     A     a      0        2   0.9
1   1     A     b      2        0   0.8
2   1     A     c      1        2   0.7
3   2     D     d      3        1   0.8
4   2     D     e      4        2   0.9

groupby操作后

each_pid = data.groupby(data['pid']).apply(lambda g: g.sort_values(['predict', 'prob'], ascending=False))

我有一個像這樣的 dataframe。

      pid query title  label  predict  prob  score
pid                                               
1   0   1     A     a      0        2   0.9      6
    2   1     A     c      1        2   0.7     18
    1   1     A     b      2        0   0.8     42
2   4   2     D     e      4        2   0.9    186
    3   2     D     d      3        1   0.8     90

現在我想要另一個列名稱weight ，其值根據每個組的順序如下。

      pid query title  label  predict  prob  score  weight
pid                                               
1   0   1     A     a      0        2   0.9      6    1
    2   1     A     c      1        2   0.7     18    2
    1   1     A     b      2        0   0.8     42    3
2   4   2     D     e      4        2   0.9    186    1
    3   2     D     d      3        1   0.8     90    2

以及如何獲得每組的前2行

      pid query title  label  predict  prob  score  weight
pid                                               
1   0   1     A     a      0        2   0.9      6    1
    2   1     A     c      1        2   0.7     18    2
2   4   2     D     e      4        2   0.9    186    1
    3   2     D     d      3        1   0.8     90    2

任何人都可以幫忙嗎？

編輯。 感謝@Akshay Kandul 和@Allen 的幫助。 該代碼作為

data = data.groupby(level=0).head(2)
data['weight'] = data.groupby(level=0).cumcount()+1

Answer 1

您可以在索引級別為零的分組上使用 cumcount function

each_pid['weight'] = each_pid.groupby(level=[0]).cumcount()

output

        label  pid  predict  prob query title  weight
pid                                                 
1   0      0    1        2   0.9     A     a       0
    2      1    1        2   0.7     A     c       1
    1      2    1        0   0.8     A     b       2
2   4      4    2        2   0.9     D     e       0
    3      3    2        1   0.8     D     d       1

如果您真的希望它以 1 開頭，那么只需執行以下操作：

each_pid['weight'] = each_pid.groupby(level=[0]).cumcount()+1

output

        label  pid  predict  prob query title  weight
pid                                                 
1   0      0    1        2   0.9     A     a       1
    2      1    1        2   0.7     A     c       2
    1      2    1        0   0.8     A     b       3
2   4      4    2        2   0.9     D     e       1
    3      3    2        1   0.8     D     d       2

如果您需要每組的前 2 條記錄，您可以使用以下代碼：

print(each_pid.groupby(level=0).head(2))

Answer 2

根據你已經得到的 each_pid，你可以使用 apply 來計算順序：

each_pid['weight'] = \
(
     each_pid.groupby(level=0)
         .apply(lambda x: pd.Series(range(1,len(x)+1)))
         .values
)

each_pid
Out[263]: 
       pid query title  label  predict  prob  weight
pid                                                 
1   0    1     A     a      0        2   0.9       1
    2    1     A     c      1        2   0.7       2
    1    1     A     b      2        0   0.8       3
2   4    2     D     e      4        2   0.9       1
    3    2     D     d      3        1   0.8       2

如果你需要每組的前 2 名，你可以這樣做：

each_pid.groupby(level=0).head(2)
Out[273]: 
       pid query title  label  predict  prob  weight
pid                                                 
1   0    1     A     a      0        2   0.9       1
    2    1     A     c      1        2   0.7       2
2   4    2     D     e      4        2   0.9       1
    3    2     D     d      3        1   0.8       2