Numpy: sum condicional

Tengo la siguiente matriz numpy:

import numpy as np arr = np.array([[1,2,3,4,2000], [5,6,7,8,2000], [9,0,1,2,2001], [3,4,5,6,2001], [7,8,9,0,2002], [1,2,3,4,2002], [5,6,7,8,2003], [9,0,1,2,2003] ]) 

Entiendo np.sum(arr, axis=0) para proporcionar el resultado:

 array([ 40, 28, 36, 34, 16012]) 

Lo que me gustaría hacer ( sin un bucle for ) es sumr las columnas en función del valor de la última columna para que el resultado provisto sea:

 array([[ 6, 8, 10, 12, 4000], [ 12, 4, 6, 8, 4002], [ 8, 10, 12, 4, 4004], [ 14, 6, 8, 10, 4006]]) 

Me doy cuenta de que puede ser un tramo sin hacer un bucle, pero esperando lo mejor …

Si se debe usar un bucle for, ¿cómo funcionaría?

    Intenté np.sum(arr[:, 4]==2000, axis=0) (donde sustituiría 2000 con la variable del bucle for), sin embargo, dio un resultado de 2

    Puede hacer esto en número puro usando una aplicación inteligente de np.diff y np.add.reduceat . np.diff le dará los índices donde cambia la columna de la derecha:

     d = np.diff(arr[:, -1]) 

    np.where convertirá su índice booleano d en los índices enteros que np.add.reduceat espera:

     d = np.where(d)[0] 

    reduceat también esperará ver un índice de cero, y todo debe ser cambiado en uno:

     indices = np.r_[0, e + 1] 

    Usar np.r_ aquí es un poco más conveniente que np.concatenate porque permite escalar. La sum entonces se convierte en:

     result = np.add.reduceat(arr, indices, axis=0) 

    Esto se puede combinar en una sola línea, por supuesto:

     >>> result = np.add.reduceat(arr, np.r_[0, np.where(np.diff(arr[:, -1]))[0] + 1], axis=0) >>> result array([[ 6, 8, 10, 12, 4000], [ 12, 4, 6, 8, 4002], [ 8, 10, 12, 4, 4004], [ 14, 6, 8, 10, 4006]]) 

    Estoy publicando una solución simple con pandas y una con itertools

     import pandas as pd df = pd.DataFrame(arr) x = df.groupby(4).sum().reset_index()[range(5)] #range(5) adjusts ordering x[4] *= 2 np.array(x) array([[ 6, 8, 10, 12, 4000], [ 12, 4, 6, 8, 4002], [ 8, 10, 12, 4, 4004], [ 14, 6, 8, 10, 4006]]) 

    También puedes usar itertools

     np.array([sum(x[1]) for x in itertools.groupby(arr, key = lambda k: k[-1])]) array([[ 6, 8, 10, 12, 4000], [ 12, 4, 6, 8, 4002], [ 8, 10, 12, 4, 4004], [ 14, 6, 8, 10, 4006]]) 

    Enfoque # 1: NumPy basado en reducción de sum

    Aquí hay uno basado en np.add.reduceat

     def groupbycol(a, assume_sorted_col=False, colID=-1): if assume_sorted_col==0: # If a is not already sorted by that col, use argsort indices for # that colID and re-arrange rows accordingly sidx = a[:,colID].argsort() a_s = a[sidx] # sorted by colID col of input array else: a_s = a # Get group shifting indices cut_idx = np.flatnonzero(np.r_[True, a_s[1:,colID] != a_s[:-1,colID]]) # Use those indices to setup sum reduction at intervals along first axis return np.add.reduceat(a_s, cut_idx, axis=0) 

    Ejecución de la muestra

     In [64]: arr Out[64]: array([[ 1, 2, 3, 4, 2000], [ 5, 6, 7, 8, 2000], [ 9, 0, 1, 2, 2001], [ 3, 4, 5, 6, 2001], [ 7, 8, 9, 0, 2002], [ 1, 2, 3, 4, 2002], [ 5, 6, 7, 8, 2003], [ 9, 0, 1, 2, 2003]]) In [65]: # Shuffle rows off input array to create a generic last col (not sorted) ...: np.random.seed(0) ...: np.random.shuffle(arr) In [66]: arr Out[66]: array([[ 5, 6, 7, 8, 2003], [ 9, 0, 1, 2, 2001], [ 5, 6, 7, 8, 2000], [ 9, 0, 1, 2, 2003], [ 3, 4, 5, 6, 2001], [ 1, 2, 3, 4, 2000], [ 1, 2, 3, 4, 2002], [ 7, 8, 9, 0, 2002]]) In [67]: groupbycol(arr, assume_sorted_col=False, colID=-1) Out[67]: array([[ 6, 8, 10, 12, 4000], [ 12, 4, 6, 8, 4002], [ 8, 10, 12, 4, 4004], [ 14, 6, 8, 10, 4006]]) 

    Enfoque # 2: Aprovechar la multiplicación de matrices

    Básicamente, podríamos reemplazar ese np.add.reduceat con una creación de máscara de difusión + multiplicación de matriz, por lo tanto, aprovechar el BLAS rápido y que también funciona para una columna genérica no ordenada:

     import pandas as pd def groupbycol_matmul(a, colID=-1): mask = pd.Series(a[:,colID]).unique()[:,None] == arr[:,colID] return mask.dot(arr) 

    Es posible que desee echar un vistazo a numpy_indexed . Con ello podrás hacer:

     import numpy as np import numpy_indexed as npi arr = np.array([[1,2,3,4,2000], [5,6,7,8,2000], [9,0,1,2,2001], [3,4,5,6,2001], [7,8,9,0,2002], [1,2,3,4,2002], [5,6,7,8,2003], [9,0,1,2,2003] ]) result = npi.GroupBy(arr[:, 4]).sum(arr)[1] >>>[[ 6 8 10 12 4000] [ 12 4 6 8 4002] [ 8 10 12 4 4004] [ 14 6 8 10 4006]]