Os pandas obtêm os mais altos n registros dentro de cada grupo

161

Suponha que eu tenha o DataFrame dos pandas assim:

>>> df = pd.DataFrame({'id':[1,1,1,2,2,2,2,3,4],'value':[1,2,3,1,2,3,4,1,1]})
>>> df
   id  value
0   1      1
1   1      2
2   1      3
3   2      1
4   2      2
5   2      3
6   2      4
7   3      1
8   4      1

Quero obter um novo DataFrame com os 2 principais registros para cada ID, assim:

   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

Eu posso fazer isso com a numeração de registros dentro de grupo após grupo:

>>> dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()
>>> dfN
   id  level_1  index  value
0   1        0      0      1
1   1        1      1      2
2   1        2      2      3
3   2        0      3      1
4   2        1      4      2
5   2        2      5      3
6   2        3      6      4
7   3        0      7      1
8   4        0      8      1
>>> dfN[dfN['level_1'] <= 1][['id', 'value']]
   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

Mas existe uma abordagem mais eficaz / elegante para fazer isso? E também existe uma abordagem mais elegante para numerar registros dentro de cada grupo (como a função da janela SQL row_number () ).

Roman Pekar
fonte
1
"top-n" não significa "as n primeiras / primeiras / linhas da cabeça", como você está procurando! Significa "as n linhas com os maiores valores".
smci 13/05/19

Respostas:

180

Você tentou df.groupby('id').head(2)

Saída gerada:

>>> df.groupby('id').head(2)
       id  value
id             
1  0   1      1
   1   1      2 
2  3   2      1
   4   2      2
3  7   3      1
4  8   4      1

Lembre-se de que talvez você precise solicitar / classificar antes, dependendo dos seus dados)

EDIT: Como mencionado pelo questionador, use df.groupby('id').head(2).reset_index(drop=True)para remover o multíndice e achatar os resultados.

>>> df.groupby('id').head(2).reset_index(drop=True)
    id  value
0   1      1
1   1      2
2   2      1
3   2      2
4   3      1
5   4      1
dorvak
fonte
1
Sim, acho que é isso. Negligenciou isso de alguma forma. Você conhece uma boa maneira de numerar registros dentro do grupo?
Roman Pekar
4
Para se ter necessidade de saída eu, eu também acrescentou.reset_index(drop=True)
Roman Pekar
1
O github.com/pydata/pandas/pull/5510 acabou de ser incorporado; estará em 0,13, novo método para fazer exatamente isso chamada cumcount(número os registros em cada grupo)
Jeff
1
@Jeff boas notícias. Eu gostaria de ter mais tempo para contribuir para Pandas :(
Roman Pekar
3
Para tornar @dorvak a resposta mais completa, se você quiser os 2 menores valores id, faça df.sort_values(['id', 'value'], axis=0).groupby('id').head(2). Outro exemplo, o maior valor por idé dado por df.sort_values(['id', 'value'], axis=0).groupby('id').tail(1).
Elmex80s
129

Desde 0.14.1 , agora você pode fazer nlargeste nsmallestem um groupbyobjeto:

In [23]: df.groupby('id')['value'].nlargest(2)
Out[23]: 
id   
1   2    3
    1    2
2   6    4
    5    3
3   7    1
4   8    1
dtype: int64

Há uma ligeira estranheza que você obtenha o índice original, em lá também, mas isso pode ser realmente útil dependendo do que o seu índice de originais foi .

Se você não estiver interessado, pode fazer isso .reset_index(level=1, drop=True)para se livrar completamente.

(Nota: A partir da 0.17.1, você também poderá fazer isso em um DataFrameGroupBy, mas por enquanto ele só funciona com Seriese SeriesGroupBy.)

LondonRob
fonte
Existe uma maneira de obter unique_limit(n)? Como eu quero os primeiros n valores únicos? Se eu pedir para nlargestele irá classificar toda a df que pode ser caro
citynorman
2
Isso não funciona nos casos em que você faz uma agregação no groupby? Por exemplo, df.groupby([pd.Grouper(freq='M'), 'A'])['B'].count().nlargest(5, 'B') isso apenas retorna os 5 principais em toda a série, não por cada grupo
geominded em 21/06/19
A afirmação de que isso agora também é possível em DataFrameGroupBys parece ser falsa, a solicitação pull vinculada parece adicionar apenas a s nlargestsimples DataFrame. O que é bastante lamentável, porque e se você quiser selecionar mais de uma coluna?
Oulenz # 23/19
7

Às vezes, classificar todos os dados com antecedência leva muito tempo. Podemos agrupar primeiro e fazer topk para cada grupo:

g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)
Chaffee Chen
fonte