Gấu trúc có được hồ sơ n hàng đầu trong mỗi nhóm

Giả sử tôi có DataFrame như thế này:

>>> df = pd.DataFrame({'id':[1,1,1,2,2,2,2,3,4],'value':[1,2,3,1,2,3,4,1,1]})
>>> df
   id  value
0   1      1
1   1      2
2   1      3
3   2      1
4   2      2
5   2      3
6   2      4
7   3      1
8   4      1

Tôi muốn nhận một DataFrame mới với 2 bản ghi hàng đầu cho mỗi id, như thế này:

   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

Tôi có thể làm điều đó với các bản ghi đánh số trong nhóm sau nhóm bằng cách:

>>> dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()
>>> dfN
   id  level_1  index  value
0   1        0      0      1
1   1        1      1      2
2   1        2      2      3
3   2        0      3      1
4   2        1      4      2
5   2        2      5      3
6   2        3      6      4
7   3        0      7      1
8   4        0      8      1
>>> dfN[dfN['level_1'] <= 1][['id', 'value']]
   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

Nhưng có cách tiếp cận hiệu quả / thanh lịch hơn để làm điều này? Và cũng có cách tiếp cận thanh lịch hơn đối với các bản ghi số trong mỗi nhóm (như hàm cửa sổ SQL row_number () ).

Bạn đã thử chưa df.groupby('id').head(2)

Ouput được tạo ra:

>>> df.groupby('id').head(2)
       id  value
id             
1  0   1      1
   1   1      2 
2  3   2      1
   4   2      2
3  7   3      1
4  8   4      1

(Hãy nhớ rằng bạn có thể cần phải đặt hàng / sắp xếp trước, tùy thuộc vào dữ liệu của bạn)

EDIT: Như người hỏi đã đề cập, sử dụng df.groupby('id').head(2).reset_index(drop=True)để loại bỏ multindex và làm phẳng kết quả.

>>> df.groupby('id').head(2).reset_index(drop=True)
    id  value
0   1      1
1   1      2
2   2      1
3   2      2
4   3      1
5   4      1

Kể từ 0.14.1 , bây giờ bạn có thể thực hiện nlargestvà nsmallesttrên một groupbyđối tượng:

In [23]: df.groupby('id')['value'].nlargest(2)
Out[23]: 
id   
1   2    3
    1    2
2   6    4
    5    3
3   7    1
4   8    1
dtype: int64

Có một điều kỳ lạ là bạn cũng có được chỉ số gốc trong đó, nhưng điều này có thể thực sự hữu ích tùy thuộc vào chỉ số ban đầu của bạn là gì .

Nếu bạn không quan tâm đến nó, bạn có thể làm .reset_index(level=1, drop=True)để loại bỏ nó hoàn toàn.

(Lưu ý: Từ 0.17.1, bạn cũng có thể thực hiện việc này trên DataFramegroupBy nhưng hiện tại nó chỉ hoạt động với Seriesvà SeriesGroupBy.)

Đôi khi việc sắp xếp toàn bộ dữ liệu phía trước rất tốn thời gian. Chúng tôi có thể nhóm trước và thực hiện topk cho mỗi nhóm:

g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)