如何選擇熊貓每組的前3行？

Question

我得到一個數據幀熊貓這樣的：

id prob 
0 1 0.5 
1 1 0.6 
2 1 0.4 
3 1 0.2 
4 2 0.3 
6 2 0.5 
... 

我想這組由「ID」，排序降序排列，並獲得第3的概率各組。請注意，某些羣體包含的行小於3 最後，我希望得到一個二維數組，如：

[[1, 0.6, 0.5, 0.4], [2, [0.5, 0.3]]...] 

我如何能做到這一點有熊貓嗎？ 謝謝！

Answer 1

使用sort_values，groupby和head：

df.sort_values(by=['id','prob'], ascending=[True,False]).groupby('id').head(3).values 

輸出：

array([[ 1. , 0.6], 
     [ 1. , 0.5], 
     [ 1. , 0.4], 
     [ 2. , 0.5], 
     [ 2. , 0.3]]) 

繼@COLDSPEED鉛：

df.sort_values(by=['id','prob'], ascending=[True,False])\ 
    .groupby('id').agg(lambda x: x.head(3).tolist())\ 
    .reset_index().values.tolist() 

輸出：

[[1, [0.6, 0.5, 0.4]], [2, [0.5, 0.3]]] 

Answer 2

您可以使用GROUPBY和nlargest

df.groupby('id').prob.nlargest(3).reset_index(1,drop = True) 

id 
1 0.6 
1 0.5 
1 0.4 
2 0.5 
2 0.3 

對於數組

df1 = df.groupby('id').prob.nlargest(3).unstack(1)#.reset_index(1,drop = True)#.set_index('id') 
np.column_stack((df1.index.values, df1.values)) 

你得到

array([[ 1. , 0.5, 0.6, 0.4, nan, nan], 
     [ 2. , nan, nan, nan, 0.3, 0.5]]) 

Answer 3

如果您正在尋找陣列列的數據幀，你可以使用np.sort：

df = df.groupby('id').prob.apply(lambda x: np.sort(x.values)[:-4:-1]) 
df 

id 
1 [0.6, 0.5, 0.4] 
2   [0.5, 0.3] 

來獲取值，reset_index和訪問：

df.reset_index().values 

array([[1, array([ 0.6, 0.5, 0.4])], 
     [2, array([ 0.5, 0.3])]], dtype=object) 

Answer 4

[[n, g.nlargest(3).tolist()] for n, g in df.groupby('id').prob] 

[[1, [0.6, 0.5, 0.4]], [2, [0.5, 0.3]]]