pandas groupby aplicar en múltiples columnas para generar una nueva columna

Me gusta generar una nueva columna en el dataframe de pandas usando groupby-apply.

Por ejemplo, tengo un dataframe:

df = pd.DataFrame({'A':[1,2,3,4],'B':['A','B','A','B'],'C':[0,0,1,1]}) 

e intente generar una nueva columna ‘D’ por groupby-apply.

Esto funciona:

 df = df.assign(D=df.groupby('B').C.apply(lambda x: x - x.mean())) 

como (creo) devuelve una serie con el mismo índice con el dataframe:

 In [4]: df.groupby('B').C.apply(lambda x: x - x.mean()) Out[4]: 0 -0.5 1 -0.5 2 0.5 3 0.5 Name: C, dtype: float64 

Pero si trato de generar una nueva columna utilizando varias columnas, no puedo asignarla directamente a una nueva columna. Así que esto no funciona:

  df.assign(D=df.groupby('B').apply(lambda x: xA - xCmean())) 

regresando

 TypeError: incompatible index of inserted column with frame index 

Y de hecho, el grupo por aplicación devuelve:

 In [8]: df.groupby('B').apply(lambda x: xA - xCmean()) Out[8]: B A 0 0.5 2 2.5 B 1 1.5 3 3.5 Name: A, dtype: float64 

Yo podría hacer

 df.groupby('B').apply(lambda x: xA - xCmean()).reset_index(level=0,drop=True)) 

pero parece detallado y no estoy seguro de si esto funcionará como se espera siempre.

Entonces, mi pregunta es: (i) ¿cuándo los grupos de pandas por solicitud devuelven una serie con un índice similar frente a una serie de múltiples índices? (ii) ¿hay una mejor manera de asignar una nueva columna por grupo aplicando a múltiples columnas?

group_keys=False en el groupby

 df.assign(D=df.groupby('B', group_keys=False).apply(lambda x: xA - xCmean())) 

Salida:

  ABCD 0 1 A 0 0.5 1 2 B 0 1.5 2 3 A 1 2.5 3 4 B 1 3.5 

Para este caso no creo que sea necesario incluir la columna A en la aplicación, podemos usar transform

 df.A-df.groupby('B').C.transform('mean') Out[272]: 0 0.5 1 1.5 2 2.5 3 3.5 dtype: float64 

Y puedes volver a asignarlo.

 df['diff']= df.A-df.groupby('B').C.transform('mean') df Out[274]: ABC diff 0 1 A 0 0.5 1 2 B 0 1.5 2 3 A 1 2.5 3 4 B 1 3.5