Tengo un problema en el que necesito tomar grupos de filas de un marco de datos donde el número de elementos en un grupo excede un cierto número (límite). Para esos grupos, necesito tomar algunas filas de la cabeza y la fila de la cola.

Estoy usando el siguiente código

train = train[train.groupby('id').id.transform(len) > headRows]
groups = pd.concat([train.groupby('id').head(headRows),train.groupby('id').tail(1)]).sort_index()

Esto funciona. Pero la primera línea, es muy lenta :(. 30 minutos o más.

¿Hay alguna forma de hacer que la primera línea sea más rápida? Si no uso la primera línea, hay índices duplicados del resultado de la segunda línea, lo que arruina las cosas.

Gracias de antemano Saludos

Nota: El marco de datos de mi tren tiene alrededor de 70,000 grupos de diferentes tamaños de grupo en alrededor de 700,000 filas. De hecho, se desprende de mi otra pregunta como se puede ver aquíProcesamiento de datos con agregar columnas dinámicamente en el marco de datos de Python Pandas. Jeff dio una gran respuesta allí, pero falla si el tamaño del grupo es menor o igual al parámetro que paso en la cabeza (parámetro) al concatenar mis filas como en la respuesta de Jeff: En [31]: grupos = concat .....