如何用Numpy将函数应用于两个向量生成二维矩阵（避免循环）

需求说明 #

给定两个向量：

a = [1,2,3]
b = [5,6,7]

以及运算函数：

def calc(i, j):
    return (i+j)/2

需要将函数应用于两个向量的所有元素对，生成目标矩阵，且避免使用列表推导和for循环，优先使用Numpy实现。

解决方案 #

方法1：利用Numpy广播机制（高效推荐）

Numpy的广播机制可自动扩展数组维度，实现无循环的元素对运算：

import numpy as np

# 将列表转为Numpy数组
a_arr = np.array([1, 2, 3])
b_arr = np.array([5, 6, 7])

# 将a转为列向量，与b进行广播运算
result = (a_arr[:, np.newaxis] + b_arr) / 2
print(result)

输出结果：

[[3.  3.5 4. ]
 [3.5 4.  4.5]
 [4.  4.5 5. ]]

方法2：结合np.vectorize复用自定义函数

若需保留自定义的calc函数，可使用np.vectorize将其转为向量化函数，再结合广播实现：

import numpy as np

def calc(i, j):
    return (i + j) / 2

a_arr = np.array([1, 2, 3])
b_arr = np.array([5, 6, 7])

# 向量化自定义函数
vec_calc = np.vectorize(calc)
# 传入列向量a和行向量b，利用广播完成所有元素对运算
result = vec_calc(a_arr[:, np.newaxis], b_arr)
print(result)

输出结果与方法1一致。

原理说明 #

a_arr[:, np.newaxis]将原本形状为(3,)的一维数组转换为(3,1)的列向量，与形状为(3,)的b_arr运算时，Numpy会自动将b_arr扩展为(1,3)的行向量，两者相加后得到(3,3)的矩阵，最终除以2得到目标结果，全程无需手动循环。

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