Pandas фильтровать, группировать и затем преобразовывать

Question

У меня есть pandas фрейм данных, который выглядит следующим образом:

  df =
      a          b
      a1.        1
      a2         0
      a1         0
      a3         1
      a2         1
      a1         1

Я хотел бы сначала отфильтровать b на 1, а затем сгруппировать по a и подсчитать количество раз происходит каждая группа (назовите этот столбец count), а затем присоедините этот столбец с исходным df. b гарантированно будет иметь хотя бы один раз 1 для каждого значения a.

Ожидаемый результат:

  df =
      a          b.    count
      a1.        1      2
      a2         0.     1
      a1         0.     2
      a3         1      1
      a2         1.     1
      a1         1      2

Я пробовал:

 df['count] = df.groupby('a').b.transform('size')

Но при этом учитываются и нули. Я хочу сначала отфильтровать b == 1.

Я также пробовал:

df['count'] = df[df['b' == 1].groupby('a').b.transform('size')

Но это вводит nans в столбце счетчика?

Как я могу сделать это одной строкой?

Answer 1

Проверьте и получите условие, применимое к b, затем sum

df['b'].eq(1).groupby(df['a']).transform('sum')
Out[103]: 
0    2.0
1    1.0
2    2.0
3    1.0
4    1.0
5    2.0
Name: b, dtype: float64