Pandas: asigna un índice a cada grupo identificado por groupby

Al usar groupby (), ¿cómo puedo crear un DataFrame con una nueva columna que contenga un índice del número de grupo, similar a dplyr::group_indices en R. Por ejemplo, si tengo

 >>> df=pd.DataFrame({'a':[1,1,1,2,2,2],'b':[1,1,2,1,1,2]}) >>> df ab 0 1 1 1 1 1 2 1 2 3 2 1 4 2 1 5 2 2 

¿Cómo puedo obtener un DataFrame como

  ab idx 0 1 1 1 1 1 1 1 2 1 2 2 3 2 1 3 4 2 1 3 5 2 2 4 

(el orden de los índices idx no importa)

Aquí hay una forma concisa de usar drop_duplicates y merge para obtener un identificador único.

 group_vars = ['a','b'] df.merge( df.drop_duplicates( group_vars ).reset_index(), on=group_vars ) ab index 0 1 1 0 1 1 1 0 2 1 2 2 3 2 1 3 4 2 1 3 5 2 2 5 

El identificador en este caso va 0,2,3,5 (solo un residuo del índice original), pero esto podría cambiarse fácilmente a 0,1,2,3 con un reset_index(drop=True) adicional reset_index(drop=True) .

Una forma simple de hacerlo sería concatenar sus columnas de agrupación (de modo que cada combinación de sus valores represente un elemento singularmente distinto), luego convertirlo en una categoría de pandas y mantener solo sus tags:

 df['idx'] = pd.Categorical(df['a'].astype(str) + '_' + df['b'].astype(str)).codes df ab idx 0 1 1 0 1 1 1 0 2 1 2 1 3 2 1 2 4 2 1 2 5 2 2 3 

Edición: se modificaron labels propiedades de las labels a los codes ya que las anteriores parecen estar en desuso

Edit2: Agregó un separador como lo sugirió Authman Apatira

Aquí está la solución utilizando ngroup de un comentario anterior de Constantino , para aquellos que aún buscan esta función (el equivalente de dplyr::group_indices en R, si estuviera intentando dplyr::group_indices en Google con esas palabras clave como yo). Esto también es aproximadamente un 25% más rápido que la solución dada por maxliving según mi propio tiempo.

 >>> import pandas as pd >>> df = pd.DataFrame({'a':[1,1,1,2,2,2],'b':[1,1,2,1,1,2]}) >>> df['idx'] = df.groupby(['a', 'b']).ngroup() >>> df ab idx 0 1 1 0 1 1 1 0 2 1 2 1 3 2 1 2 4 2 1 2 5 2 2 3 >>> %timeit df['idx'] = create_index_usingduplicated(df, grouping_cols=['a', 'b']) 1.83 ms ± 67.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) >>> %timeit df['idx'] = df.groupby(['a', 'b']).ngroup() 1.38 ms ± 30 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) 

Una forma en la que creo que es más rápida que la respuesta aceptada actual en aproximadamente un orden de magnitud (resultados de tiempo a continuación):

 def create_index_usingduplicated(df, grouping_cols=['a', 'b']): df.sort_values(grouping_cols, inplace=True) # You could do the following three lines in one, I just thought # this would be clearer as an explanation of what's going on: duplicated = df.duplicated(subset=grouping_cols, keep='first') new_group = ~duplicated return new_group.cumsum() 

Resultados de tiempo:

 a = np.random.randint(0, 1000, size=int(1e5)) b = np.random.randint(0, 1000, size=int(1e5)) df = pd.DataFrame({'a': a, 'b': b}) In [6]: %timeit df['idx'] = pd.Categorical(df['a'].astype(str) + df['b'].astype(str)).codes 1 loop, best of 3: 375 ms per loop In [7]: %timeit df['idx'] = create_index_usingduplicated(df, grouping_cols=['a', 'b']) 100 loops, best of 3: 17.7 ms per loop 

No estoy seguro de que esto sea un problema tan trivial. Aquí hay una solución un tanto complicada que primero ordena las columnas de agrupación y luego verifica si cada fila es diferente de la fila anterior y, si es así, se acumula por 1. Ver más abajo para obtener una respuesta con datos de cadena.

 df.sort_values(['a', 'b']).diff().fillna(0).ne(0).any(1).cumsum().add(1) 

Salida

 0 1 1 1 2 2 3 3 4 3 5 4 dtype: int64 

Así que, dividiendo esto en pasos, veamos la salida de df.sort_values(['a', 'b']).diff().fillna(0) que verifica si cada fila es diferente de la fila anterior. Cualquier entrada que no sea cero indica un nuevo grupo.

  ab 0 0.0 0.0 1 0.0 0.0 2 0.0 1.0 3 1.0 -1.0 4 0.0 0.0 5 0.0 1.0 

Un nuevo grupo solo necesita tener una única columna diferente, de modo que esto es lo que .ne(0).any(1) verifica – no es igual a 0 para ninguna de las columnas. Y luego solo una sum acumulada para hacer un seguimiento de los grupos.

Respuesta para columnas como cadenas

 #create fake data and sort it df=pd.DataFrame({'a':list('aabbaccdc'),'b':list('aabaacddd')}) df1 = df.sort_values(['a', 'b']) 

salida de df1

  ab 0 aa 1 aa 4 aa 3 ba 2 bb 5 cc 6 cd 8 cd 7 dd 

Adopte un enfoque similar al verificar si el grupo ha cambiado

 df1.ne(df1.shift().bfill()).any(1).cumsum().add(1) 0 1 1 1 4 1 3 2 2 3 5 4 6 5 8 5 7 6 

Definitivamente no es la solución más sencilla, pero esto es lo que haría (comentarios en el código):

 df=pd.DataFrame({'a':[1,1,1,2,2,2],'b':[1,1,2,1,1,2]}) #create a dummy grouper id by just joining desired rows df["idx"] = df[["a","b"]].astype(str).apply(lambda x: "".join(x),axis=1) print df 

Eso generaría un idx único para cada combinación de a y b .

  ab idx 0 1 1 11 1 1 1 11 2 1 2 12 3 2 1 21 4 2 1 21 5 2 2 22 

Pero este sigue siendo un índice bastante tonto (piense en algunos valores más complejos en las columnas a y b . Así que borremos el índice:

 # create a dictionary of dummy group_ids and their index-wise representation dict_idx = dict(enumerate(set(df["idx"]))) # switch keys and values, so you can use dict in .replace method dict_idx = {y:x for x,y in dict_idx.iteritems()} #replace values with the generated dict df["idx"].replace(dict_idx,inplace=True) print df 

Eso produciría la salida deseada:

  ab idx 0 1 1 0 1 1 1 0 2 1 2 1 3 2 1 2 4 2 1 2 5 2 2 3