如何为DataFrame各分组保留不同预设行数？

假设你有如下示例DataFrame：

import pandas as pd

data = {
    'city': ['Buenos Aires', 'Paris', 'Barcelona', 'Buenos Aires', 'Buenos Aires', 'Paris', 'Barcelona'],
    'number': [14, 23, 12, 14, 14, 25, 13]
}
df = pd.DataFrame(data)

要实现按city分组后，给不同城市保留指定行数，可按以下步骤操作：

1. 定义各城市的行数限制
   创建字典，明确每个目标城市需要保留的行数，其他城市可选择保留全部数据或自定义规则：

row_limits = {
    'Buenos Aires': 2,  # 保留2行
    'Paris': 1          # 保留1行
}

2. 分组后按规则截取数据
   使用groupby结合apply，对每个分组判断是否在行数限制字典中，执行对应截取逻辑：

def filter_group(group):
    city = group.name
    if city in row_limits:
        # 取前N行，若要随机选取可改用group.sample(n=row_limits[city])
        return group.head(row_limits[city])
    else:
        # 其他城市保留全部数据，也可改为自定义规则（比如保留0行）
        return group

filtered_df = df.groupby('city', group_keys=False).apply(filter_group)

执行后得到的filtered_df结果示例：

city  number
0  Buenos Aires      14
3  Buenos Aires      14
1         Paris      23
2     Barcelona      12
6     Barcelona      13

补充说明 #

• 如果需要随机选取指定行数而非取前N行，将group.head(row_limits[city])替换为group.sample(n=row_limits[city], random_state=42)即可（random_state用于固定随机结果）。
• 若要对其他城市也设置规则，直接在row_limits字典中添加对应键值对即可。

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