Numpy高效的矩阵向量差的逐元素argmin

在本文中，我们将介绍如何使用NumPy高效地计算矩阵向量差的逐元素argmin。在解决某些机器学习、数据分析和科学计算问题时，我们经常需要计算矩阵向量差的逐元素argmin。举个例子，假设我们有一个矩阵A和一个向量b，我们想要找到每一行矩阵A和向量b之间差值的最小值所在的列号。这种问题在人脸识别、声音信号处理和自然语言处理等领域中都是非常常见的。

阅读更多：Numpy 教程

问题描述

设矩阵 $A$ 的行数为 $M$ ，列数为 $N$ ；向量 $b$ 的长度为 $N$ 。我们需要寻找每一行 $A_i$ 和向量 $b$ 之间差的最小值所在的列号 $j$ ，即：

$j = \text{argmin}( |A_{i,j} – b_j| ), \quad 1\le i\le M$

暴力解法

最简单的思路就是暴力遍历每一行的每一个元素，这样的时间复杂度为 $\mathcal{O}(MN)$ 。在小规模的矩阵中没有问题，但当我们处理大规模矩阵时会变得非常低效。

import numpy as np
M, N = 1000, 1000
A = np.random.rand(M, N)
b = np.random.rand(N)

res = []
for i in range(M):
    d = np.abs(A[i,:] - b)
    res.append(np.argmin(d))

NumPy的解法

NumPy提供了一种简单且高效的解法，使用np.argmin()函数和广播机制即可。如下所示，我们可以使用np.argmin()函数和广播机制计算矩阵 $A$ 和向量 $b$ 之间的差值，最后返回每个差值中的最小值对应的列号。

res = np.argmin(np.abs(A-b[:,np.newaxis]), axis=1)

代码的解释如下：

b[:,np.newaxis]将向量 $b$ 转换成一个列向量，即矩阵 $b’$ ；
np.abs(A-b[:,np.newaxis])计算矩阵 $A$ 和 $b’$ 之间的差值，并且取绝对值；
np.argmin(np.abs(A-b[:,np.newaxis]), axis=1)返回每行差值中的最小值对应的列号。

这种方法的时间复杂度为 $\mathcal{O}(MN)$ 。

加速

我们可以进一步优化上述算法，通过对矩阵转置或者重新交换 $A$ 和 $b$ 来避免高速缓存(cache)未命中的情况。具体来说，当 $M>N$ 时，我们可以将 $A$ 转置后按列进行计算；反之，当 $M\le N$ 时，我们可以交换 $A$ 和 $b$ ，然后对 $A$ 的每一行计算逐元素argmin。如下所示：

if M > N:
    res = np.argmin(np.abs(A-b[:,np.newaxis].T), axis=0)
else:
    res = np.zeros(M, dtype=int)
    for i in range(M):
        res[i] = np.argmin(np.abs(A[i,:] - b))