python - Pandas 在每个组中获得最高的不同记录

Question

在这种情况下，我想带来每个 id 的最高值，但数量不同。
也就是说，我正在寻找 'id'=1 的 5 个最高值，'id'=2 的 3 个最高值，等等。
我有这个代码，它只为我带来每组固定数量的值。

import random

df = pd.DataFrame({'id':[1,1,1,1,1,1,1,1,2,2,2,2,3,3,3,3,4,4,4]})
df['value'] = np.random.randint(0, 99, df.shape[0])
df.groupby(['id']).apply(lambda x: x.nlargest(2,['value'])).reset_index(drop=True)

id = 1 --> 5
id = 2 --> 3
id = 3 --> 2
id = 4 --> 2

Answer 1

IUC:

def my_largest(d):
    # define a dictionary with the specific
    # number of largest rows to grab for
    # each `'id'`
    nlim = {1: 5, 2: 3, 3: 2, 4: 2}

    # When passing a dataframe from a
    # `groupby` to the callable used in
    # the `apply`, Pandas will attach an
    # attribute `name` to that dataframe
    # whose value is the disctint group
    # the dataframe represents.  In this
    # case, that will be the `'id'` because
    # we grouped by `'id'`
    k = nlim[d.name]
    return d.nlargest(k, ['value'])

df.groupby('id').apply(my_largest).reset_index(drop=True)

    id  value
0    1     96
1    1     83
2    1     58
3    1     49
4    1     43
5    2     66
6    2     40
7    2     33
8    3     90
9    3     54
10   4     83
11   4     23

同样的事情，但具有更通用的功能

现在这个函数可以接受任何规范字典。此外，我还包含了一个参数以在存在 'id' 的情况下使用默认值。规范字典中不存在的。

def my_largest(d, nlrg_dict, nlrg_dflt=5, **kw):
    k = nlrg_dict.get(d.name, nlrg_dflt)
    return d.nlargest(k, **kw)

现在，您可以看到我们在函数外定义了字典......

nlim = {1: 5, 2: 3, 3: 2, 4: 2}

...并通过 apply 将其传递给函数

df.groupby('id').apply(
    my_largest, nlrg_dict=nlim, columns=['value']
).reset_index(drop=True)