python之使用索引列表的列表对 numpy 矩阵的行进行求和

使用索引列表和应用函数对 numpy 数组进行切片，是否可以矢量化(或非矢量化方式)？矢量化将是大型矩阵的理想选择

import numpy as np 
index = [[1,3], [2,4,5]] 
a = np.array( 
       [[ 3,  4,  6,  3], 
        [ 0,  1,  2,  3], 
        [ 4,  5,  6,  7], 
        [ 8,  9, 10, 11], 
        [12, 13, 14, 15], 
        [1, 1,    4,  5]]) 

对 index 中的行索引组求和，给出:

np.array([[8, 10, 12, 14], 
          [17, 19, 24, 37]]) 

请您参考如下方法:

方法 #1:这是一种几乎*矢量化的方法 -

def sumrowsby_index(a, index): 
    index_arr = np.concatenate(index) 
    lens = np.array([len(i) for i in index]) 
    cut_idx = np.concatenate(([0], lens[:-1].cumsum() )) 
    return np.add.reduceat(a[index_arr], cut_idx) 

*几乎是因为使用循环理解计算 lens 的步骤，但由于我们只是获取长度并且那里不涉及任何计算，因此该步骤不会影响任何时间大路。

sample 运行-

In [716]: a 
Out[716]:  
array([[ 3,  4,  6,  3], 
       [ 0,  1,  2,  3], 
       [ 4,  5,  6,  7], 
       [ 8,  9, 10, 11], 
       [12, 13, 14, 15], 
       [ 1,  1,  4,  5]]) 
 
In [717]: index 
Out[717]: [[1, 3], [2, 4, 5]] 
 
In [718]: sumrowsby_index(a, index) 
Out[718]:  
array([[ 8, 10, 12, 14], 
       [17, 19, 24, 27]]) 

方法#2:我们可以利用 numpy.dot 的快速矩阵乘法执行这些总和减少，为我们提供了另一种方法，如下所列 -

def sumrowsby_index_v2(a, index): 
    lens = np.array([len(i) for i in index]) 
    id_ar = np.zeros((len(lens), a.shape[0])) 
    c = np.concatenate(index) 
    r = np.repeat(np.arange(len(index)), lens)     
    id_ar[r,c] = 1 
    return id_ar.dot(a)