我有一个2D矩阵,我需要对矩阵元素的一个子集求和,给出两个索引列表imp_list
和bath_list
。这就是我现在正在做的事情:
s = 0.0
for i in imp_list:
for j in bath_list:
s += K[i,j]
似乎很慢。什么是更好的解决方案来执行总和?
答案 0 :(得分:4)
如果您正在使用大型阵列,那么通过使用NumPy自己的索引程序而不是Python的for
循环,您应该可以大大提高速度。
在一般情况下,您可以使用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