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

【问题标题】：Partial sum over an array given a list of indices给定索引列表的数组的部分和
【发布时间】：2015-09-25 06:16:22
【问题描述】：

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

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

这似乎很慢。执行求和的更好解决方案是什么？

【问题讨论】：

它很慢，因为它是一个 n^2 算法。就是这样。您必须以一种或另一种方式访问所有这些职位。
imp_list 有多大，bath_list 有多大？
在这里使用 numpy 会很有帮助。这是链接docs.scipy.org/doc/numpy/reference/generated/…

标签： python arrays performance numpy multidimensional-array

【解决方案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

【讨论】：

谢谢，正是我所需要的，它为我节省了很多时间，而且我还可以在其他功能中重复使用您的方法。