能否用纯NumPy API实现不同尺寸数组间的相似度计算？

给定以下脚本：

import numpy as np
from numpy.linalg import norm

a = np.array([(1, 2, 3), (1, 4, 9), (2, 4, 4)])
b = np.array([(1, 3, 3), (1, 5, 9)])
r = sum([min(norm(a-e, ord=1, axis=1)) for e in b])

这段代码计算了不同尺寸NumPy数组a和b之间的相似度r。能否完全通过NumPy API实现该计算以提升效率？

可以通过NumPy的向量化操作完全替代原代码中的Python循环，实现更高效的计算。核心是利用广播机制一次性计算所有元素对的L1距离，再完成后续的最小距离提取与求和。

实现代码 #

import numpy as np

a = np.array([(1, 2, 3), (1, 4, 9), (2, 4, 4)])
b = np.array([(1, 3, 3), (1, 5, 9)])

# 广播计算所有b元素与a元素的L1距离，结果形状为(len(b), len(a))
distances = np.abs(a[np.newaxis, :, :] - b[:, np.newaxis, :]).sum(axis=2)
# 对每个b元素对应的距离数组取最小值，再求和得到最终结果
r = distances.min(axis=1).sum()

效率优势 #

• 原代码依赖Python列表推导式循环，每次循环调用norm都会带来Python解释器的额外开销，数据量越大，性能越差。
• 向量化实现完全在NumPy的底层C环境中执行，避免了循环的解释器开销，对于大规模数组，性能提升会非常明显。

内容的提问来源于stack exchange，提问作者Paul Jurczak