DataFrame: добавить столбец с размером группы

У меня есть следующий фрейм данных:

    fsq digits  digits_type
0    1   1       odd
1    2   1       odd
2    3   1       odd
3    11  2       even
4    22  2       even
5    101 3       odd
6    111 3       odd

и я хочу добавить последний столбец count, содержащий номер fsq, принадлежащего группе цифр, т.е.

    fsq digits  digits_type   count
0    1   1       odd          3
1    2   1       odd          3
2    3   1       odd          3
3    11  2       even         2
4    22  2       even         2
5    101 3       odd          2
6    111 3       odd          2

Так как есть 3 строки fsq с цифрами, равными 1, 2 строки fsq с цифрами, равными 2, и т.д.

Ответ 1

In [395]: df['count'] = df.groupby('digits')['fsq'].transform(len)

In [396]: df
Out[396]: 
   fsq  digits digits_type  count
0    1       1         odd      3
1    2       1         odd      3
2    3       1         odd      3
3   11       2        even      2
4   22       2        even      2
5  101       3         odd      2
6  111       3         odd      2

[7 rows x 4 columns]

Ответ 2

В общем, вы должны использовать методы, определенные Pandas, где это возможно. Это часто будет более эффективным.

В этом случае вы можете использовать 'size' в том же духе, что и df.groupby('digits')['fsq'].size():

df = pd.concat([df]*10000)

%timeit df.groupby('digits')['fsq'].transform('size')  # 3.44 ms per loop
%timeit df.groupby('digits')['fsq'].transform(len)     # 11.6 ms per loop