Question

这段代码在Python中很慢。如何使用numpy优化它。数组m和z已经是numpy数组，但我想我需要以某种方式对其进行矢量化，尽管我找不到矢量化函数也可以访问值坐标的示例。

for i in xrange(states):
    for j in xrange(states):
        if i != j and z[i, j] != 0 and m[i, j] < 0.0:
            m[i, i] += m[i, j]
            m[i, j] = 0.0
            z[i, j] = 0

Answer 1

您可以使用索引数组将代码转换为矢量化Numpy：

import numpy as np

i, j = np.indices([states, states])
bool_index = (i != j) & (z != 0) & (m < 0.0)
z[bool_index] = 0

但是既然你已经知道i != j会评估什么，那么使用对角阵列会更快。这也使得在m上执行稍微笨拙的操作变得更容易：

I = np.eye(states, dtype=bool)
bool_index = ~I & (z != 0) & (m < 0.0)

m[I] += (bool_index * m).sum(axis=1)
m[bool_index] = 0.0
z[bool_index] = 0

请注意，&运算符执行np.bitwise_and而非np.logical_and，但在这种情况下它们是等效的，因为所有操作数都是布尔数组。

如何使用numpy在数组上优化这个双循环？

1 个答案: