如何在pandas数据框中保留每组的前3行?

时间:2017-02-16 03:42:41

标签: python pandas grouping

我有一个如下所示的列表:

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中执行此操作?

2 个答案:

答案 0 :(得分:8)

set_indexgroupbySeriesGroupBy.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

答案 1 :(得分:5)

如果 count 列已按降序排序,那么您只需使用groupby.head从每组中获取前三行:

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

enter image description here

否则,您可以按var1对数据框进行分组,并使用nlargest检索前3位计数的三行:

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

enter image description here