CODEWITHSUNDEEP
pythonpandasgroup-by
Raymont
Raymont

Reputation: 293

Pandas get topmost distinct records within each group

in this case I want to bring the highest values of each id but in different quantities. That is, I'm looking for the 5 highest values for the 'id'=1, the 3 highest values for the 'id'=2, etc. I have this code that only brings me a fixed amount of values per group.

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

Upvotes: 2

Views: 196

Answers (1)

piRSquared
piRSquared

Reputation: 294506

IIUC:

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

Same thing but with a more generalized function

Now this function can take any specification dictionary. Also, I've included a parameter to use a default in the case that there is an 'id' that doesn't exist in the specification dictionary.

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

Now, you can see we define the dictionary outside the function ...

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

... and pass it to the function via apply

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

Upvotes: 2

Related Questions