NumPy矩阵操作

Question

我希望能够计算所有 i 和 j 的以下数值：

M_ki = Sum[A_ij - A_ik - A_kj + A_kk, 1 <= j <= n]

如何在不使用显式循环的情况下，使用Numpy（Python）来完成它？

谢谢！

- yassin

Answer 1

以下是解决此类问题的一般策略。

首先，编写一个小脚本，其中循环在两个不同的函数中显式编写，并在末尾进行测试，确保两个函数完全相同：

import numpy as np
from numpy import newaxis

def explicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    for k in range(n):
        for i in range(n):
            for j in range(n):
                m[k,i] += a[i,j] - a[i,k] - a[k,j] + a[k,k]
    return m

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    for k in range(n):
        for i in range(n):
            for j in range(n):
                m[k,i] += a[i,j] - a[i,k] - a[k,j] + a[k,k]
    return m

a = np.random.randn(10,10)
assert np.allclose(explicit(a), implicit(a), atol=1e-10, rtol=0.)

接下来，通过逐步编辑implicit函数来向量化它，每一步都要运行脚本以确保它们继续保持不变：

第一步

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    for k in range(n):
        for i in range(n):
            m[k,i] = (a[i,:] - a[k,:]).sum() - n*a[i,k] + n*a[k,k]
    return m

第二步

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    m = - n*a.T + n*np.diag(a)[:,newaxis]
    for k in range(n):
        for i in range(n):
            m[k,i] += (a[i,:] - a[k,:]).sum()
    return m

第三步

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    m = - n*a.T + n*np.diag(a)[:,newaxis]
    m += (a.T[newaxis,...] - a[...,newaxis]).sum(1)
    return m

看，最后一个版本都没有循环了。要将这种类型的方程向量化，广播是正确的方法！

警告：确保explicit是你想要向量化的方程。我不确定那些不依赖于j的术语是否也应该被求和。

- pberkes

---

非常好的答案，而且是很棒的通用建议！ - Joe Kington

非常感谢。对未来的建议很棒 :) - yassin