Como adicionar uma coluna a um dataframe agrupado no Pandas

Na análise de dados, muitas vezes é necessário agrupar dados e realizar cálculos em cada grupo. O Pandas oferece uma maneira conveniente de fazer isso por meio de sua função groupby. Uma tarefa comum é contar os valores de uma coluna dentro de cada grupo e adicionar uma coluna contendo essas contagens ao dataframe.

Considere o dataframe df:

df = pd.DataFrame({'c':[1,1,1,2,2,2,2],'type':['m','n','o','m','m','n','n']})

Para contar os valores do tipo para cada c, podemos usar a função value_counts no dataframe agrupado:

g = df.groupby('c')['type'].value_counts().reset_index(name='t')

Isso cria um novo dataframe g com as contagens de grupo. Para adicionar uma coluna a g com o tamanho de cada grupo, podemos usar a função de transformação:

g['size'] = df.groupby('c')['type'].transform('size')

transform aplica uma função a cada grupo no dataframe original e retorna uma série com seu índice alinhado ao dataframe original. Neste caso, usamos a função size para contar o número de elementos em cada grupo e atribuí-lo ao novo tamanho da coluna. O dataframe g resultante agora ficará assim:

   c type  t  size
0  1    m  1     3
1  1    n  1     3
2  1    o  1     3
3  2    m  2     4
4  2    n  2     4

Isso demonstra uma maneira direta de adicionar uma nova coluna a um dataframe agrupado com base nos resultados de uma agregação groupby.