Pandas 有效地获取每个数据帧组中具有多个 n 值的前 n 行

从像

df = pd.DataFrame(
    [
        {'name': 'a', 'id': 0, 'category': '1'},
        {'name': 'b', 'id': 1, 'category': '1'},
        {'name': 'c', 'id': 2, 'category': '1'},
        {'name': 'd', 'id': 3, 'category': '1'},
        {'name': 'e', 'id': 4, 'category': '2'},
        {'name': 'f', 'id': 5, 'category': '2'},
        {'name': 'g', 'id': 6, 'category': '2'},
        {'name': 'h', 'id': 7, 'category': '3'},
        {'name': 'i', 'id': 8, 'category': '3'},
        {'name': 'j', 'id': 9, 'category': '3'},
    ]
)

或

  name  id category
0    a   0        1
1    b   1        1
2    c   2        1
3    d   3        1
4    e   4        2
5    f   5        2
6    g   6        2
7    h   7        3
8    i   8        3
9    j   9        3

这样的数据帧开始，我想提取每个类别的第一 n 行，每个类别都有不同的 n 如果 n 对于所有类别都是相同的，那么 这个问题的解决方案 将会很好而且快速：

df_head = df.groupby(by=['category'], as_index=False, sort=False).head(2)

给出

  name  id category
0    a   0        1
1    b   1        1
4    e   4        2
5    f   5        2
7    h   7        3
8    i   8        3

但是如果我想要不同的 n ，我就只剩下像 this

num_grabs = {'1': 3, '2': 1, '3': 2}
df_apply = df.groupby(
    by=['category'],
    as_index=False,
    sort=False
).apply(
    lambda x: x.head(num_grabs[x.loc[x.index[0], 'category']])
).reset_index(drop=True)

这样的方法可以正确给出

  name  id category
0    a   0        1
1    b   1        1
2    c   2        1
3    e   4        2
4    h   7        3
5    i   8        3

但是，即使对于中等大小的数据来说，这也非常慢。有没有更有效的方法来解决多 n 问题？

虽然给出的 apply 方法有效，但它对于大型数据集来说效率可能很低，因为它循环遍历每个组。以下是一种可能更快的替代方法，它使用 NumPy 和一些聪明的索引：

import pandas as pd
import numpy as np

def get_top_n_per_group(df, group_col, n_dict):
    # 获取组大小和累积计数
    group_sizes = df[group_col].value_counts(sort=False).to_numpy()
    group_cumsum = np.cumsum(group_sizes)

    # 创建一个布尔掩码，指示要保留的行
    n_values = np.array([n_dict[key] for key in sorted(n_dict)])
    max_n = n_values.max()
    row_indices = np.arange(len(df))
    group_starts = np.concatenate(([0], group_cumsum[:-1]))
    within_group_indices = row_indices - group_starts.repeat(group_sizes)
    mask = within_group_indices < n_values.repeat(group_sizes)

    # 应用掩码并返回过滤后的 DataFrame
    return df[mask]

# 示例用法
df = pd.DataFrame({
    'name': 'a b c d e f g h i j'.split(),
    'id': range(10),
    'category': '1 1 1 1 2 2 2 3 3 3'.split()
})

num_grabs = {'1': 3, '2': 1, '3': 2}
df_head = get_top_n_per_group(df, 'category', num_grabs)
print(df_head)

这段代码的工作原理如下：

1. 计算组大小和累积计数： 这有助于我们稍后确定每组的起始位置。
2. 创建布尔掩码：
3. 我们创建一个 row_indices 数组来表示 DataFrame 中的每一行。
4. group_starts 数组存储每组的起始索引。
5. within_group_indices 计算每行相对于其组的起始位置。
6. 最后， mask 通过检查 within_group_indices 是否小于每组所需的 n 值来标识要保留的行。
7. 应用掩码并返回过滤后的 DataFrame： 我们简单地将掩码应用于原始 DataFrame 以获取所需的输出。

此方法避免了显式循环，并且应该比基于 apply 的方法快得多，尤其是在大型数据集上。

此解决方案假定“类别”列已排序，或者至少每组的顺序不影响结果。如果的数据集中“类别”列的值的顺序不同，则需要对数据进行排序或调整解决方案以首先收集类别值。