如何在pandas數據框中保留每組的前3行？

Question

我有一個如下所示的列表：

var1     var2    count
 A        abc      4
 A        abc      3
 A        abc      2
 A        abc      1
 A        abc      1
 B        abc      7
 B        abc      5
 B        abc      2
 B        abc      1
 B        abc      1
 C        abc      4
 C        abc      3
 C        abc      2
 C        abc      1
 C        abc      1

 ....

我想創建一個新的數據框，其中包含每組的前3個“計數”結果。 它應該如下所示：

     var1     var2    count
      A        abc      4
      A        abc      3
      A        abc      2
      B        abc      7
      B        abc      5
      B        abc      2
      C        abc      4
      C        abc      3
      C        abc      2
      ....

有沒有一種方便的方法在使用head（）在Python中執行此操作？

Answer 1

使用set_index ， groupby和SeriesGroupBy.nlargest解決方案：

df = df.set_index('var2').groupby("var1")['count'].nlargest(3).reset_index()
print (df)
  var1 var2  count
0    A  abc      4
1    A  abc      3
2    A  abc      2
3    B  abc      7
4    B  abc      5
5    B  abc      2
6    C  abc      4
7    C  abc      3
8    C  abc      2

Answer 2

如果count列已按降序排序，那么您只需使用groupby.head從每個組中獲取前三行：

df.groupby("var1").head(3)

否則，您可以按var1數據框進行分組，並使用nlargest檢索前3個計數的三行：

df.groupby("var1", group_keys=False).apply(lambda g: g.nlargest(3, "count"))