如何为每个组保留预定义的不同行数？

我有一个包含几列的数据框。我想按城市订购并保留 'city' == 'Buenos Aires' 一定数量的行。对于 'city' == 'Paris' ，还有另一个行数。有没有办法做到这一点？ 这里 显示了一种为每个组保持相同行数的方法，但我想要不同的行数。

    city            number
0   Buenos Aires    14
1   Paris           23
2   Barcelona       12
3   Buenos Aires    14
4   Buenos Aires    14
... ...             ...

可以使用 Pandas 的 groupby 和 head 函数的组合，为每个组保留不同数量的行。以下是操作方法：

1. 创建一个字典来存储每个城市的所需行数。 例如：

rows_per_city = {
    'Buenos Aires': 2,  # 为 Buenos Aires 保留 2 行
    'Paris': 3,         # 为 Paris 保留 3 行
    # 为其他城市添加更多条目
}

1. 使用 groupby 按城市对 DataFrame 进行分组。

2. 使用 head 函数获取每个组中的所需行数。 可以使用字典查找来获取每个城市的正确行数。

import pandas as pd

# 示例 DataFrame
df = pd.DataFrame({
    'city': ['Buenos Aires', 'Paris', 'Barcelona', 'Buenos Aires', 'Buenos Aires', 'Paris', 'Paris', 'Paris', 'Barcelona'],
    'number': [14, 23, 12, 14, 14, 23, 12, 14, 14]
})

# 为每个城市定义所需的行数
rows_per_city = {
    'Buenos Aires': 2,
    'Paris': 3,
}

# 使用 lambda 函数对每个组应用 head 并从字典中获取 n 值
df_filtered = df.groupby('city').apply(lambda x: x.head(rows_per_city.get(x.name, len(x))))

# 如果索引级别过多，则重置索引
df_filtered = df_filtered.reset_index(level=0, drop=True)

# 打印结果
print(df_filtered)

这将创建一个新的 DataFrame，其中 Buenos Aires 的前两行和 Paris 的前三行，以及其他城市的所有行（因为它们未在字典中指定）。

以下是对代码的解释：

• groupby('city') 按“城市”列对 DataFrame 进行分组。
• apply(lambda x: ...) 将一个函数应用于每个组。
• x 代表一个组（在本例中是一个 DataFrame）。
• x.head(...) 获取组中的前 n 行。
• rows_per_city.get(x.name, len(x)) 获取当前城市 ( x.name ) 的所需行数。如果未找到城市，则使用 len(x) （组中的总行数），从而保留所有行。
• reset_index(level=0, drop=True) 删除由 groupby 操作创建的多余索引级别。

此方法允许通过为所需的每个组在 rows_per_city 字典中指定所需的组和行数来控制每个组保留的行数。