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]])

按索引中的行索引组求和，得出：

使用列表理解。。。 对于索引中的每个索引列表，创建一个新列表，该列表是这些索引中某个索引中的行列表。从这里，我们有一个numpy数组列表，可以应用sum方法。在numpy数组中，sum将从添加的数组中返回每个元素的新数组，该数组将为您提供所需的内容：

np.array([sum([a[r] for r in i]) for i in index])

给予：

array([[ 8, 10, 12, 14],
       [17, 19, 24, 27]])

使用列表理解。。。 对于索引中的每个索引列表，创建一个新列表，该列表是这些索引中某个索引中的行列表。从这里，我们有一个numpy数组列表，可以应用sum方法。在numpy数组中，sum将从添加的数组中返回每个元素的新数组，该数组将为您提供所需的内容：

np.array([sum([a[r] for r in i]) for i in index])

给予：

array([[ 8, 10, 12, 14],
       [17, 19, 24, 27]])

方法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)

*几乎是因为计算镜头的一步是循环理解，但由于我们只是得到长度，没有计算，这一步不会对计时产生任何大的影响

样本运行-

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：我们可以利用快速矩阵乘法来执行这些求和运算，这为我们提供了另一种方法，如下所示-

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)

方法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)

*几乎是因为计算镜头的一步是循环理解，但由于我们只是得到长度，没有计算，这一步不会对计时产生任何大的影响

样本运行-

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：我们可以利用快速矩阵乘法来执行这些求和运算，这为我们提供了另一种方法，如下所示-

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)