All rows of First N items of a group of data in dataset based on another column in pandas

Viewed 16

Let's consider I have this dataset:

   name comp   item      type
    A    c1    item21     t1
    A    c1    item231    t1
    A    c1    item3      t1
    B    c3    item23     t1
    B    c3    item1      t1
    B    c3    p3251      t1
    C    c4    item1      t1
    C    c4    p32sd      t1
    C    c4    item512    t1
    D    c5    item242    t2
    D    c5    item1      t2
    F    c6    item4      t2
    F    c6    item24     t2
    H    c7    item4125   t2
    H    c7    item3      t2
    H    c7    item14     t2
    K    c8    item1      t2
    K    c8    p3223      t2

I want to select all items of first n [names,comp] of each type:

For example all items of first 2 names-comp of each type the expected df would be:

   name comp   item      type
    A    c1    item21     t1
    A    c1    item231    t1
    A    c1    item3      t1
    B    c3    item23     t1
    B    c3    item1      t1
    B    c3    p3251      t1
    D    c5    item242    t2
    D    c5    item1      t2
    F    c6    item4      t2
    F    c6    item24     t2

Does anybody have any idea how to do this?

1 Answers

Try this:

cols = ['type', 'name', 'comp']

# The first 2 name-comp of each type
tmp = df[cols].drop_duplicates().groupby('type').head(2)

# All rows that match the criteria
result = tmp.merge(df, left_on=cols, right_on=cols)

If you want no intermediary data frame:

df[cols].drop_duplicates().groupby('type').head(2).merge(df, left_on=cols, right_on=cols)
Related