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

我得到这样的数据帧。

 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功能零

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

输出

 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开始，然后只是做这个：如果你需要前2记录每个组的

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

输出

 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 

您可以使用此代码：

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

Answer 2

基于each_pid你已经有了，您可以使用适用于计算顺序：如果你需要每个小组的前2

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 

，你可以这样做：

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