python - 使 np.sums 的 for 循环运行得更快的方法

Question

我有这段代码，它迭代矩阵列表，获取它们的 np.sum 并对变量进行加/减。这段代码是我正在处理的一个更大的模拟问题的一部分，由于每次运行整个程序时这段小代码都会运行大约 10mi+ 次，所以我试图使其尽可能高效。有什么办法可以简化这个过程，使其运行得更快吗？我寻找了一些矢量化技巧，但没有发现如何提高效率

#mi could be of any size
mi = np.random.randint(2,size=(4,4))
#the arrays in list_of_blocks could be of any size so the np.sum of the list would end in a ValueError
list_of_blocks = [np.array([[0],[1]]),np.array([[0,0,0],[0,1,0]])]
nee = np.sum(mi)
n1 =0
mrxnr= 0
for block in list_of_blocks:
    nee -= np.sum(block)
    n1 += np.sum(block)
    mrxnr += block.shape[0]*block.shape[1]

Answer 1

矢量化的工作原理是对 numpy 数组应用一些操作，这些数组的数据存储在连续的内存块中。列表的情况并非如此，因此循环对于这种情况似乎没问题。

但是，您的代码中有一些内容需要优化。list_of_blocks 中的数组总是 0 或 1 吗？如果是这种情况，count_nonzero 比 sum 快得多。

此外，block.shape[0]*block.shape[1] 相当于 block.size，速度也更快。

最后，nee 和 n1 是相关的，因此不需要计算 sum(block) 两次 对于每个循环周期。优化后的代码:

nee = np.sum(mi)
n1 =0
mrxnr= 0
for block in list_of_blocks:
    n1 += np.count_nonzero(block)
    mrxnr += block.size
nee -=n1