2d numpy，有效地分配非零索引 [row,col] [row,row] 和 [col,col] 的最小值

Question

我正在使用矩阵来计算图中节点之间的关系组合（细节无关紧要）。

我有一个 N*N 邻接矩阵，行和列对应节点。因此，位置 [5,7] 是节点 5 与节点 7 的次数。[7,5] 也是如此。位置[3,3]是节点3总共出现了多少次，所以总共出现了多少次。

在每个循环中，我都必须减少我的矩阵。我取一个大小为 n 的向量，然后用该向量减去矩阵对角线。所以，我正在减少每个节点的总数：就像我的矩阵中的 [1,1] 和 [2,2] 和 [3,3] 等等。

希望到目前为止我说得通。这是这个问题。

此时，我已经修改了矩阵的对角线。现在，我想修改每个位置 [i,j]

i != j and matrix[i,j] != 0  

我想修改它：

matrix[i,j] = min(matrix[i,i],matrix[j][j])

现在我当然可以遍历每个索引对 (i,j) 并执行我刚才写的操作。但这很慢。我希望有一些聪明的数学或 numpy 技巧可以使这更快。

谢谢！

Answer 1

首先，在进行任何优化之前，您应该分析一下：对于在程序的整个生命周期中只需要几十毫秒的事情，或者只占总运行时间的一小部分，试图聪明是没有意义的。

也就是说，您可以利用广播来矢量化取最小值：

def slow(arr):
    out = arr.copy()
    for (i, j), x in np.ndenumerate(arr):
        if i != j and arr[i,j] != 0:
            out[i,j] = min(arr[i, i], arr[j, j])
    return out

def fast(arr):
    diag = arr.diagonal()
    mins = np.minimum(diag, diag[:, None])
    out = np.where(arr != 0, mins, arr)
    out[np.diag_indices_from(arr)] = diag
    return out

这给了我

In [61]: a = np.random.randint(0, 10, (100, 100))

In [62]: (slow(a) == fast(a)).all()
Out[62]: True

In [63]: %timeit slow(a)
11.9 ms ± 188 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [64]: %timeit fast(a)
62.8 µs ± 916 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)