python pandas新增加一列_python-如何在 pandas 中`groupby`之后根据顺序添加一列？

我的数据就像,我想计算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

经过分组操作

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

现在,我想要另一列名称权重,该权重的值取决于每个组的顺序,如下所示.

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

有人可以帮忙吗？

data = data.groupby(level=0).head(2)

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