不同维度numpy数组从末尾反向对位相加的实现方法

原代码错误原因 #

• 循环过程中直接修改了原始b数组的取值，后续迭代用到的b已经不是初始输入的b
• 错误将单条a的行和整个b数组做加法，触发numpy广播后得到的结果不符合预期

方案1：修改原有循环逻辑 #

如果你需要保留循环写法，可以调整为如下代码：

import numpy as np
c = list()
a = np.array([[1,2,3],[1,2,3],[1,2,3],[4,5,6],[7,8,9]])
b = np.array([[4,5,6],[5,8,9]])
n = len(b)
for i in range(n):
    # 分别取a和b的倒数第i+1个元素相加
    res = a[-(i+1)] + b[-(i+1)]
    c.append(res)
# 翻转列表得到要求的输出顺序
c = np.array(c[::-1])
print(c)

运行后输出：

[[ 8 10 12]
 [12 16 18]]

方案2：向量化实现（推荐） #

numpy原生支持数组对位运算，无需循环即可实现需求，性能更高：

import numpy as np
a = np.array([[1,2,3],[1,2,3],[1,2,3],[4,5,6],[7,8,9]])
b = np.array([[4,5,6],[5,8,9]])
# 直接取a的最后len(b)行与b对位相加即可
c = a[-len(b):] + b
print(c)

运行结果和上述循环写法完全一致。

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