Pandas：如何限制透视表中城市索引的行数最多为3行

原始数据 #

首先定义目标DataFrame：

import pandas as pd

df = pd.DataFrame({
'city': ['New York', 'New York', 'New York', 'New York', 'New York', 'New York', 'Los Angeles', 'Los Angeles', 'Houston', 'Houston', 'Houston', 'Boston', 'Boston', 'Boston', 'Boston'],
'airport': ['LGA', 'EWR', 'JFK', 'TEB', 'CWD', 'TTN', 'LAX', 'BUR', 'IAH', 'HOU', 'EFD', 'BOS', 'ACK', 'MVY', 'WST'],
'distance': [38, 32, 8, 78, 120, 180, 8, 19, 90, 78, 120, 9, 97, 72, 150]
})

原有实现及问题 #

通过以下代码生成排序后的透视表：

pivot_table = pd.pivot_table(df, index = ['city', 'airport'], values = 'distance')

sorted_table = pivot_table.reset_index().sort_values(['city', 'distance'], ascending=[1,0]).set_index(['city', 'airport'])

此时部分城市（如波士顿、纽约）关联的机场数量超过3个，需要限制每个city对应的结果行数最多为3行，目标输出如下：

distance
       city    airport  
     Boston        WST           150
                   ACK            97
                   MVY            72
    Houston        EFD           120
                   IAH            90
                   HOU            78
Los Angeles        BUR            19
                   LAX             8
   New York        TTN           180
                   CWD           120
                   TEB            78

解决方案 #

在排序完成后，通过groupby结合head(3)实现每组最多保留3行数据：

# 生成透视表并重置索引、排序
pivot_table = pd.pivot_table(df, index=['city', 'airport'], values='distance')
sorted_df = pivot_table.reset_index().sort_values(['city', 'distance'], ascending=[True, False])

# 按city分组，每组取前3行，再重新设置多级索引
limited_table = sorted_df.groupby('city').head(3).set_index(['city', 'airport'])

说明 #

• 先完成排序：按city升序、distance降序排列，确保每个城市下的机场按距离从远到近排序
• groupby('city').head(3)会对每个城市的分组保留前3条记录，刚好筛选出每个城市距离最远的3个机场
• 最后通过set_index(['city', 'airport'])恢复原有的多级索引格式，得到目标透视表

内容的提问来源于stack exchange，提问作者equanimity