给定索引列表的数组的部分和

Question

我有一个二维矩阵，我需要对矩阵元素的一个子集求和，给定两个索引列表imp_list和bath_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())。

下面是一些数据来说明这些改进:

K = np.arange(1000000).reshape(1000, 1000)
imp_list = range(100)  # [0, 1, 2, ..., 99]
bath_list = range(200) # [0, 1, 2, ..., 199]
def method1():
    s = 0
    for i in imp_list:
        for j in bath_list:
            s += K[i,j]
    return s
def method2():
    return K[np.ix_(imp_list, bath_list)].sum()
def method3():
    return K[:100, :200].sum()

:

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