Pandas получить самые выдающиеся записи в каждой группе

Question

в этом случае я хочу привести самые высокие значения каждого идентификатора, но в разных количествах. То есть я ищу 5 самых высоких значений для 'id' = 1, 3 самых высоких значения для 'id' = 2, et c. У меня есть этот код, который приносит мне только фиксированное количество значений на группу.

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

IIU C:

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)