首页 > 解决方案 > `np.add.at` 到二维数组

问题描述

我正在寻找np.add.at().

预期的行为如下。

augend = np.zeros((10, 10))
indices_for_dim0 = np.array([1, 5, 2])
indices_for_dim1 = np.array([5, 3, 1])
addend = np.array([1, 2, 3])

### some procedure substituting np.add.at ###

assert augend[1, 5] == 1
assert augend[5, 3] == 2
assert augend[2, 1] == 3

任何建议都会有所帮助!

标签: pythonnumpy

解决方案


您可以np.add.at按原样多维使用。该indices参数在描述中包含以下内容:

...如果第一个操作数有多个维度,则索引可以是数组的元组,如索引对象或切片

所以:

augend = np.zeros((10, 10))
indices_for_dim0 = np.array([1, 5, 2])
indices_for_dim1 = np.array([5, 3, 1])
addend = np.array([1, 2, 3])
np.add.at(augend, (indices_for_dim0, indices_for_dim1), addend)

更简单地说:

augend[indices_for_dim0, indices_for_dim1] += addend

如果您真的担心多维方面并且您的被加数是一个普通的连续 C 顺序数组,您可以使用ravelandravel_multi_index在一维视图上执行操作:

indices = np.ravel_multi_index((indices_for_dim0, indices_for_dim1), augend.shape)
raveled = augend.ravel()
np.add.at(raveled, indices, addend)

推荐阅读