如何在NumPy中一步实现带偏移量的数组相加操作？

NumPy没有直接实现该特定逻辑的单一内置函数，需要通过组合操作或自定义函数来完成。

需求逻辑回顾 #

给定数组：

import numpy as np
a = np.array([1,2,3,4,5])
b = np.array([10,11,12])

需要从a的索引3位置开始与b相加，具体逻辑为：

• 将a[3:]补0至与b长度一致（即补1个0，得到[4,5,0]）
• 与b对应位置相加：[4+10,5+11,0+12] = [14,16,12]
• 将结果拼接到a[:3]后，得到最终数组[1,2,3,14,16,12]

自定义实现函数 #

以下是符合需求的自定义函数：

import numpy as np

def insert(data1, data2, offset):
    # 创建长度为偏移量+data2长度的全0数组
    c = np.zeros(len(data2) + offset)
    # 写入data1的原始内容
    c[:len(data1)] = data1
    # 从指定偏移位置开始叠加data2
    c[offset:] += data2
    return c

测试验证

运行以下代码即可得到目标结果：

a = np.array([1,2,3,4,5])
b = np.array([10,11,12])
result = insert(a, b, 3)
print(result)
# 输出：[ 1.  2.  3. 14. 16. 12.]

分步实现方式 #

如果不需要封装函数，也可以分步执行：

• 创建足够长度的全0数组，写入a的内容：

  target_length = len(b) + 3
  a_extended = np.zeros(target_length)
  a_extended[:len(a)] = a
  

• 从索引3的位置叠加b：

  a_extended[3:] += b
  

• 最终a_extended就是目标数组

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