如何通过索引数组将numpy小数组广播至大数组并执行加法

错误原因分析 #

1. A[Index, Index] += B 报错原因 #

A[Index, Index]采用的是numpy的配对式高级索引，它选取的是(Index[0],Index[0])、(Index[1],Index[1])……(Index[3],Index[3])这4个独立元素，最终得到形状为(4,)的一维数组，而B是(4,4)的二维数组，两者形状无法广播匹配，因此抛出shape mismatch错误。

2. np.add.at(A[Index], Index, B) 报错原因 #

A[Index]是从原数组A中提取出Index指定的行，得到形状为(4,9)的子数组。此时用原Index数组（包含值6）去索引这个子数组的行轴（轴0），但子数组的行轴长度仅为4（索引范围0-3），6超出有效索引范围，因此触发IndexError。

正确实现方法 #

方法1：利用广播式索引实现直接累加 #

将Index转换为列向量，与原Index组合生成二维索引网格，匹配B的形状：

import numpy as np

A = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0]])
B = np.array([[2, 4, 5, 7],
              [2, 7, 1, 3],
              [6, 1, 3, 5],
              [1, 6, 3, 2]])
Index = np.array([0, 1, 5, 6])

# 核心代码
A[Index[:, None], Index] += B

Index[:, None]将原(4,)的一维数组转换为(4,1)的列向量，与(4,)的Index组合时，numpy会自动广播为(4,4)的索引矩阵，精准定位到A中需要累加B的所有位置。

方法2：使用np.ix_生成索引网格（更直观） #

np.ix_可以直接将两个一维数组转换为可广播的索引网格：

rows, cols = np.ix_(Index, Index)
A[rows, cols] += B

这种方式更易读，rows是(4,1)的行索引数组，cols是(1,4)的列索引数组，广播后同样生成(4,4)的索引矩阵。

方法3：使用np.add.at处理重复索引场景（若有重复索引需求） #

如果Index中存在重复值，直接累加可能会因为numpy的高级索引机制导致重复赋值被覆盖，此时可以用np.add.at确保所有累加操作生效：

np.add.at(A, (Index[:, None], Index), B)

注意这里要直接对原数组A进行索引，而非子数组。

执行上述任意一种方法后，A都会变为预期的结果数组。

内容的提问来源于stack exchange，提问作者E. V. Hadzen