大型稀疏矩阵的遍历和计算

Question

问题遇到的现象和发生背景

有一个（3783x3783）的二维稀疏矩阵matrix，数值只有1、-1、0这三种，其中0占据很大一部分。我想计算的是遍历元素，并计算最后的结果。最简单且最耗时的是最开始写了一个三层for循环，计算了两个小时都还没出现一个进度条。具体代码如下：

for i in range(0, N - 2):
    for j in range(i + 1, N - 1):
        for k in range(j + 1, N):
            a = matrix[i][j]
            b = matrix[i][k]
            c = matrix[j][k]
            if ((a != 0) and (b != 0) and (c != 0)):
                d = a + b + c
                if (d == 3 or d == -1):
                    balance = balance + 1
                else:
                    unbalance = unbalance + 1

于是改为稀疏矩阵，因为稀疏矩阵中只存储了有效值，所有想着没有存储到的值就等于0。虽然用了稀疏矩阵，但是发现计算还是很忙，跑了一个小时左右也没出现一点结果。具体代码如下：

ma = csr_matrix(matrix)
for i in range(0, N - 2):
    for j in range(i + 1, N - 1):
        for k in range(j + 1, N):
            a = ma.__getitem__((i,j))
            b = ma.__getitem__((i,k))
            c = ma.__getitem__((j,k))
            if ((a != 0) and (b != 0) and (c != 0)):
                d = a + b + c
                if (d == 3 or d == -1):
                    balance = balance + 1
                else:
                    unbalance = unbalance + 1

总体要计算的就是balance和unbalance的值。怀疑是不是三层for循环的原因，但是遍历就在那儿，不好改for循环。所以想请教一下大家有什么好的方法么？感谢感谢！

Answer 1

不至于这么慢吧？简单写个demo，瞬间完成，仅供参考。

>>> import numpy as np
>>> data = np.random.randint(-1, 2, (3783,3783)) # 模拟测试数据
>>> a = data[:-2, 1:-1]
>>> b = data[:-2, 2:]
>>> c = data[1:-1, 2:]
>>> d = a + b + c
>>> d = d[(a!=0)&(b!=0)&(c!=0)]
>>> balance = d[(d==3)|(d==-1)].shape[0]
>>> unbalance = d.shape[0] - balance
>>> balance, unbalance
(2118421, 2116393)