Suma acumulativa de una matriz numpy por índice

Supongamos que tiene una serie de valores que deberán sumrse

d = [1,1,1,1,1] 

y una segunda matriz que especifica qué elementos deben sumrse

 i = [0,0,1,2,2] 

El resultado se almacenará en una nueva matriz de tamaño max(i)+1 . Entonces, por ejemplo, i=[0,0,0,0,0] sería equivalente a sumr todos los elementos de d y almacenar el resultado en la posición 0 de una nueva matriz de tamaño 1 .

Traté de implementar esto usando

 c = zeros(max(i)+1) c[i] += d 

Sin embargo, la operación += agrega cada elemento solo una vez, dando así el resultado inesperado de

     [1,1,1] 

    en lugar de

     [2,1,2] 

    ¿Cómo se podría implementar correctamente este tipo de sum?

    Esta solución debería ser más eficiente para matrices grandes (itera sobre los posibles valores de índice en lugar de las entradas individuales de i ):

     import numpy as np i = np.array([0,0,1,2,2]) d = np.array([0,1,2,3,4]) i_max = i.max() c = np.empty(i_max+1) for j in range(i_max+1): c[j] = d[i==j].sum() print c [1. 2. 7.] 

    Si entiendo la pregunta correctamente, hay una función rápida para esto (siempre que la matriz de datos sea 1d)

     >>> i = np.array([0,0,1,2,2]) >>> d = np.array([0,1,2,3,4]) >>> np.bincount(i, weights=d) array([ 1., 2., 7.]) 

    np.bincount devuelve una matriz para todos los rangos de enteros (máx. (i)), incluso si algunas cuentas son cero

     def zeros(ilen): r = [] for i in range(0,ilen): r.append(0) i_list = [0,0,1,2,2] d = [1,1,1,1,1] result = zeros(max(i_list)+1) for index in i_list: result[index]+=d[index] print result 

    El comentario de Juh_ es la solución más eficiente. Aquí está trabajando el código:

     import numpy as np import scipy.ndimage as ni i = np.array([0,0,1,2,2]) d = np.array([0,1,2,3,4]) n_indices = i.max() + 1 print ni.sum(d, i, np.arange(n_indices)) 

    En el caso general, cuando desee sumr submatrices por tags, puede usar el siguiente código

     import numpy as np from scipy.sparse import coo_matrix def labeled_sum1(x, labels): P = coo_matrix((np.ones(x.shape[0]), (labels, np.arange(len(labels))))) res = P.dot(x.reshape((x.shape[0], np.prod(x.shape[1:])))) return res.reshape((res.shape[0],) + x.shape[1:]) def labeled_sum2(x, labels): res = np.empty((np.max(labels) + 1,) + x.shape[1:], x.dtype) for i in np.ndindex(x.shape[1:]): res[(...,)+i] = np.bincount(labels, x[(...,)+i]) return res 

    El primer método utiliza la multiplicación de matrices dispersas. El segundo es la generalización de la respuesta del usuario333700. Ambos métodos tienen una velocidad comparable:

     x = np.random.randn(100000, 10, 10) labels = np.random.randint(0, 1000, 100000) %time res1 = labeled_sum1(x, labels) %time res2 = labeled_sum2(x, labels) np.all(res1 == res2) 

    Salida:

     Wall time: 73.2 ms Wall time: 68.9 ms True