python - 在 pandas DataFrame 中基于 MultiIndex 提取组的组长

Question

我找到了 this question它演示了如何根据 MultiIndex 的第一级对 DataFrame 进行排序。所以初始的 DataFrame 看起来像

l = [[1,'A',99],[1,'B',102],[1,'C',105],[1,'D',97],[2,'A',19],[2,'B',14],[2,'C',10],[2,'D',17]]
df = pd.DataFrame(l,columns = ['idx1','idx2','col1'])
df.set_index(['idx1','idx2'],inplace=True)

# assume data has been received like this...
print df

           col1
idx1 idx2      
1    A       99
     B      102
     C      105
     D       97
2    A       19
     B       14
     C       10
     D       17

可以在每个 level=1 组中排序以产生:

           col1
idx1 idx2      
1    C      105
     B      102
     A       99
     D       97

2    A       19
     D       17
     B       14
     C       10

我的问题是:如何根据第一层提取每个组的前n个元素？

例如，如果 n=2 那么结果将是:

           col1
idx1 idx2      
1    C      105
     B      102

2    A       19
     D       17

提前感谢您的考虑和回复。

Answer 1

您只需要在已排序数据帧的 groupby 上使用 head:

df.sort_values('col1', ascending=False).groupby('idx1').head(2)

输出:

           col1
idx1 idx2      
1    C      105
     B      102
2    A       19
     D       17