是否可以使用 numpy 避免此循环？

Question

假设你有两个 numpy 矩阵：

import numpy as np

n = 800
m = 16
A = np.zeros((n, 4)) #np.random.rand(n,4) if you wish
B = np.zeros(m)      #np.random.rand(m) if you wish

在我的代码中，我需要将 B 中找到的每个值减去 A 的所有值，然后计算该平方并获得为我提供该新矩阵最小值的索引。因此，我将有 m 个最小值。从这些中，我需要最低限度。为此，我现在使用这样的循环：

min_C = np.zeros(m)

for j in range(m):
    C = A[:,2]-B[j]
    C = C*C
    min_C[j] = np.where(C == np.amin(C))[0][0]

所以我对每个 j 都有一个最小值。然后，我做 min_C 的最小值。我可以避免这个循环以获得绝对最小值吗？如果可能的话，我现在无法想象如何编码。这只是为了提高我的代码的效率，我需要它，因为这个计算是用不同的 A 和 B 完成了很多次，当然。

Answer 1

利用broadcasting，您可以将其减少到以下结果min 的(m,n) 数组：

((A[:,2]-B[:,None])**2).min()

Answer 2

我假设您的示例代码中出现的固定索引 2 将保持固定，这将问题简化为 A 是一维的情况。 如果我正确理解了这个问题并查看了您的样本输出，那么您正在寻找一个索引，而不是出现在 C 本身之一中的值，这样我们就可以用绝对值替换正方形。首先，我们通过@yatu 的广播提示创建 A[i,2] 和 B[j] 的所有组合的绝对值，然后计算沿第二个轴（大小为 n）的 argmin 的最小值

c = np.abs(A[:,2]-B[:,None])
answer = np.min(np.argmin(c, axis=1))

我希望我得到你所要求的正确，如果要改变 A 的固定索引，请扩展广播：

c = np.abs(A[:,:]-B[:,None,None])
answer = np.min(np.argmin(c, axis=1))

输出将是一个 int，这与您的示例代码不同。