Pandas tiene una función de agrupación ampliamente utilizada para dividir un DataFrame en función de un mapeo correspondiente, desde el cual puede aplicar un cálculo en cada subgrupo y recombinar los resultados.
¿Se puede hacer esto de manera flexible en NumPy sin un bucle for nativo de Python? Con un bucle de Python, esto se vería así:
>>> import numpy as np >>> X = np.arange(10).reshape(5, 2) >>> groups = np.array([0, 0, 0, 1, 1]) # Split up elements (rows) of `X` based on their element wise group >>> np.array([X[groups==i].sum() for i in np.unique(groups)]) array([15, 30]) Por encima de 15 es la suma de las tres primeras filas de X y 30 es la suma de las dos restantes.
Por "flexiblemente", solo quiero decir que no nos estamos enfocando en un cálculo en particular, como la suma, el conteo, el máximo, etc., sino que pasamos cualquier cálculo a las matrices agrupadas.
Si no, ¿hay un enfoque más rápido que el anterior?
¿Qué hay de usar la matriz dispersa scipy?
import numpy as np from scipy import sparse import time x_len = 500000 g_len = 100 X = np.arange(x_len * 2).reshape(x_len, 2) groups = np.random.randint(0, g_len, x_len) # original s = time.time() a = np.array([X[groups==i].sum() for i in np.unique(groups)]) print(time.time() - s) # using scipy sparse matrix s = time.time() x_sum = X.sum(axis=1) b = np.array(sparse.coo_matrix( ( x_sum, (groups, np.arange(len(x_sum))) ), shape=(g_len, x_len) ).sum(axis=1)).ravel() print(time.time() - s) #compare print(np.abs((ab)).sum())resultado en mi pc
0.15915322303771973 0.012875080108642578 0Más de 10 veces más rápido.
¡Actualizar!
Comparemos las respuestas de @Paul Panzer y @Daniel F. Es solo un punto de referencia de suma.
import numpy as np from scipy import sparse import time # by @Daniel F def groupby_np(X, groups, axis = 0, uf = np.add, out = None, minlength = 0, identity = None): if minlength < groups.max() + 1: minlength = groups.max() + 1 if identity is None: identity = uf.identity i = list(range(X.ndim)) del i[axis] i = tuple(i) n = out is None if n: if identity is None: # fallback to loops over 0-index for identity assert np.all(np.in1d(np.arange(minlength), groups)), "No valid identity for unassinged groups" s = [slice(None)] * X.ndim for i_ in i: s[i_] = 0 out = np.array([uf.reduce(X[tuple(s)][groups == i]) for i in range(minlength)]) else: out = np.full((minlength,), identity, dtype = X.dtype) uf.at(out, groups, uf.reduce(X, i)) if n: return out x_len = 500000 g_len = 200 X = np.arange(x_len * 2).reshape(x_len, 2) groups = np.random.randint(0, g_len, x_len) print("original") s = time.time() a = np.array([X[groups==i].sum() for i in np.unique(groups)]) print(time.time() - s) print("use scipy coo matrix") s = time.time() x_sum = X.sum(axis=1) b = np.array(sparse.coo_matrix( ( x_sum, (groups, np.arange(len(x_sum))) ), shape=(g_len, x_len) ).sum(axis=1)).ravel() print(time.time() - s) #compare print(np.abs((ab)).sum()) print("use scipy csr matrix @Daniel F") s = time.time() x_sum = X.sum(axis=1) c = np.array(sparse.csr_matrix( ( x_sum, groups, np.arange(len(groups)+1) ), shape=(len(groups), g_len) ).sum(axis=0)).ravel() print(time.time() - s) #compare print(np.abs((ac)).sum()) print("use bincount @Paul Panzer @Daniel F") s = time.time() d = np.bincount(groups, X.sum(axis=1), g_len) print(time.time() - s) #compare print(np.abs((ad)).sum()) print("use ufunc @Daniel F") s = time.time() e = groupby_np(X, groups) print(time.time() - s) #compare print(np.abs((ae)).sum())SALIDA ESTÁNDAR
original 0.2882847785949707 use scipy coo matrix 0.012301445007324219 0 use scipy csr matrix @Daniel F 0.01046299934387207 0 use bincount @Paul Panzer @Daniel F 0.007468223571777344 0.0 use ufunc @Daniel F 0.04431319236755371 0El ganador es la solución bincount. Pero la solución de la matriz csr también es muy interesante.
A primera vista, la solución de matriz dispersa de @klim parecería estar vinculada a la suma. Sin embargo, podemos usarlo en el caso general al convertir entre los formatos csr y csc :
Veamos un pequeño ejemplo:
>>> m, n = 3, 8 >>> idx = np.random.randint(0, m, (n,)) >>> data = np.arange(n) >>> >>> M = sparse.csr_matrix((data, idx, np.arange(n+1)), (n, m)) >>> >>> idx array([0, 2, 2, 1, 1, 2, 2, 0]) >>> >>> M = M.tocsc() >>> >>> M.indptr, M.indices (array([0, 2, 4, 8], dtype=int32), array([0, 7, 3, 4, 1, 2, 5, 6], dtype=int32))Como podemos ver después de la conversión, la representación interna de la matriz dispersa produce los índices agrupados y ordenados:
>>> groups = np.split(M.indices, M.indptr[1:-1]) >>> groups [array([0, 7], dtype=int32), array([3, 4], dtype=int32), array([1, 2, 5, 6], dtype=int32)] >>> Podríamos haber obtenido lo mismo usando un argsort estable:
>>> np.argsort(idx, kind='mergesort') array([0, 7, 3, 4, 1, 2, 5, 6]) >>> Pero las matrices dispersas son en realidad más rápidas, incluso cuando permitimos que argsort use un algoritmo no estable más rápido:
>>> m, n = 1000, 100000 >>> idx = np.random.randint(0, m, (n,)) >>> data = np.arange(n) >>> >>> timeit('sparse.csr_matrix((data, idx, np.arange(n+1)), (n, m)).tocsc()', **kwds) 2.250748165184632 >>> timeit('np.argsort(idx)', **kwds) 5.783584725111723 Si requerimos argsort para mantener los grupos ordenados, la diferencia es aún mayor:
>>> timeit('np.argsort(idx, kind="mergesort")', **kwds) 10.507467685034499Si desea una implementación más flexible de groupby que pueda agruparse usando cualquiera de los numpy s de ufunc :
def groupby_np(X, groups, axis = 0, uf = np.add, out = None, minlength = 0, identity = None): if minlength < groups.max() + 1: minlength = groups.max() + 1 if identity is None: identity = uf.identity i = list(range(X.ndim)) del i[axis] i = tuple(i) n = out is None if n: if identity is None: # fallback to loops over 0-index for identity assert np.all(np.in1d(np.arange(minlength), groups)), "No valid identity for unassinged groups" s = [slice(None)] * X.ndim for i_ in i: s[i_] = 0 out = np.array([uf.reduce(X[tuple(s)][groups == i]) for i in range(minlength)]) else: out = np.full((minlength,), identity, dtype = X.dtype) uf.at(out, groups, uf.reduce(X, i)) if n: return out groupby_np(X, groups) array([15, 30]) groupby_np(X, groups, uf = np.multiply) array([ 0, 3024]) groupby_np(X, groups, uf = np.maximum) array([5, 9]) groupby_np(X, groups, uf = np.minimum) array([0, 6])