如何无循环快速计算两个二维NumPy数组的逐行元素差值？

我是编程/Python/NumPy的新手，若有表述不当之处敬请谅解。现有两个二维NumPy数组：

import numpy as np

a = np.array([[1,   5,              20],
              [2,   15,             float("NaN")],
              [4,   float("NaN"),   float("NaN")]])
b = np.array([[4,   float("NaN")],
              [1,   13],
              [5,   float("NaN")]])

我需要计算两个数组中对应行的所有元素间的差值，期望结果如下：

[[  3.  nan  -1.  nan -16.  nan]  # <-- b[0]与a[0]的所有元素差值
 [ -1.  11. -14.  -2.  nan  nan]  # <-- b[1]与a[1]的所有元素差值
 [  1.  nan  nan  nan  nan  nan]] # <-- b[2]与a[2]的所有元素差值

目前我已通过循环实现该需求：

outcome = []
for a_row, b_row in zip(a, b):
    outcome.append((b_row - a_row[:,None]).flatten())
outcome = np.array(outcome)

向量化实现方案 #

当然有更高效的无遍历方法，利用NumPy的广播机制就能直接对整个数组进行操作，避免Python循环的开销，代码如下：

# 向量化实现
outcome = (b[:, None, :] - a[:, :, None]).reshape(a.shape[0], -1)

代码解释 #

• b[:, None, :]：将数组b的形状从(3,2)调整为(3,1,2)，为每行添加一个中间维度，方便和a的行元素做广播运算
• a[:, :, None]：将数组a的形状从(3,3)调整为(3,3,1)，为每个元素添加一个末尾维度
• 两者相减时，NumPy会自动完成广播，生成形状为(3,3,2)的数组，其中每个位置对应b行元素与a行元素的差值
• reshape(a.shape[0], -1)：将结果的最后两个维度展平为一维，最终得到形状为(3,6)的数组，和循环实现的结果完全一致

这种方法在处理大规模数组时，性能会远优于Python循环，因为NumPy的底层操作是基于C实现的，能充分利用硬件加速。

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