Python NumPy优化需求：高效实现数组列表指定位置元素相加

问题场景 #

给定以下数组：

arr1 = [[0,0,0,0],[0,0,0,0],[0,0,0,0]]
arr2 = [[1,2],[3,4],[4,5]]
positions = [0,1,2]

需要将arr2的每个子数组插入到arr1对应子数组的positions指定起始索引处，最终得到：

[[1,2,0,0],[0,3,4,0],[0,0,4,5]]

常规循环方案在数据量大时效率偏低，以下是两种高效实现方式：

方案1：Python内置切片赋值（轻量高效） #

利用Python内置的切片赋值特性，这是底层优化过的操作，比手动逐个元素赋值快得多：

for row1, row2, pos in zip(arr1, arr2, positions):
    # 用切片替换arr1对应位置的元素区间
    row1[pos:pos+len(row2)] = row2

执行后arr1会直接被修改为目标结果。如果需要保留原arr1，可以先创建深拷贝：

import copy
arr1_copy = copy.deepcopy(arr1)
for row1, row2, pos in zip(arr1_copy, arr2, positions):
    row1[pos:pos+len(row2)] = row2

方案2：Numpy向量化操作（大规模数据首选） #

如果处理的是超大规模数组，使用Numpy的向量化操作可以彻底规避Python层面的循环，借助C底层实现大幅提升效率：

import numpy as np

# 转换为Numpy数组
np_arr1 = np.array(arr1)
np_arr2 = np.array(arr2)
pos_array = np.array(positions)

# 生成赋值所需的行、列索引
row_indices = np.repeat(np.arange(np_arr1.shape[0]), np_arr2.shape[1])
col_indices = np.concatenate([np.arange(p, p + len(row)) for p, row in zip(pos_array, np_arr2)])

# 批量完成赋值
np_arr1[row_indices, col_indices] = np_arr2.flatten()

# 可选：转换回Python列表
result = np_arr1.tolist()

常见问题说明 #

如果之前尝试切片操作失败，大概率是未正确计算切片的结束位置（pos:pos+len(row2)），或者误操作了新创建的数组而非原数组的引用。

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