如何在 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