根据索引列表对数组进行部分求和

Question

我有一个二维矩阵，我需要求和矩阵元素的子集，给定两个索引列表imp_list和{}。我现在正在做的是：

s = 0.0
for i in imp_list:
    for j in bath_list:
        s += K[i,j]

看起来很慢。有什么更好的解决方案来计算总和？在

Answer 1

如果使用大型数组，那么通过在Python的for循环上使用NumPy自己的索引例程，应该可以大大提高速度。在

在一般情况下，您可以使用np.ix_来选择要求和的矩阵的子数组：

K[np.ix_(imp_list, bath_list)].sum()

请注意，np.ix_会带来一些开销，因此如果两个列表包含连续或等间距的值，那么使用常规切片来索引数组是值得的（请参见下面的method3()）。在

以下是一些数据来说明这些改进：

^{pr2}$

然后：

In [80]: method1() == method2() == method3()
Out[80]: True

In [91]: %timeit method1()
10 loops, best of 3: 9.93 ms per loop

In [92]: %timeit method2()
1000 loops, best of 3: 884 µs per loop

In [93]: %timeit method3()
10000 loops, best of 3: 34 µs per loop