Python: добавить столбец в groupby-результат, содержащий размер каждой группы

Question

Я сгруппировал фрейм данных, чтобы получить это:

+------+----+-------+
| Type | Nr | Class |
+------+----+-------+
| One  | 01 | A1    |
| One  | 01 | A2    |
| One  | 01 | B1    |
| One  | 02 | A1    |
| One  | 02 | B1    |
| Two  | 01 | A1    |
| Two  | 01 | B1    |
| Two  | 01 | B2    |
| Two  | 02 | A1    |
+------+----+-------+

Я сделал следующее, чтобы получить число уникальных Nr значений для Типа:

DFGroup = df.groupby('Type')['Nr'].nunique().reset_index()

Это прекрасно работает:

+------+----+
| Type | Nr |
+------+----+
| One  |  2 |
| Two  |  2 |
+------+----+

, но теперь я хотел бы добавить еще один столбец к DFGroup, который содержит размер каждой группы, например:

+------+----+-------+
| Type | Nr | Count |
+------+----+-------+
| One  |  2 |     5 |
| Two  |  2 |     4 |
+------+----+-------+

Я пытался:

DFGroup['Count'] = df.groupby('Type').size()

, который дает только NAN для каждой группы.

Спасибо:)

Answer 1

Использование map:

s = df.groupby('Type').size()
DFGroup = df.groupby('Type')['Nr'].nunique().reset_index()
DFGroup['new'] = DFGroup['Type'].map(s)

print (DFGroup)
  Type  Nr  new
0  One   2    5
1  Two   2    4

Лучше использовать agg для обеих функций:

DFGroup = df.groupby('Type')['Nr'].agg([('Nr', 'nunique'),('Count','size')]).reset_index()
print (DFGroup)

  Type  Nr  Count
0  One   2      5
1  Two   2      4