使用 NumPy 对数组中的元素求和

Question

我正在用这段代码用 matplotlib 注释一个图

for position, force in np.nditer([positions, forces]):
    plt.annotate(
        "%.3g N" % force,
        xy=(position, 3),
        xycoords='data',
        xytext=(position, 2.5),
        textcoords='data',
        horizontalalignment='center',
        arrowprops=dict(arrowstyle="->")
    )

它工作正常。但是如果我在同一位置有元素，它会将多个箭头相互堆叠，即如果我有 positions = [1,1,4,4] 和 forces = [4,5,8,9]，它将在位置 1 和两个箭头在位置 4，在彼此之上。相反，我想对力求和，只在位置 1 处创建一个箭头，力为 4+5=9，在位置 4 处创建一个箭头，力为 8+9=17.

如何使用 Python 和 NumPy 执行此操作？

编辑

我想可能是这样的

import numpy as np

positions = np.array([1,1,4,4])
forces = np.array([4,5,8,9])

new_positions = np.unique(positions)
new_forces = np.zeros(new_positions.shape)

for position, force in np.nditer([positions, forces]):
    pass

Answer 1

我不确定 numpy 是否提供帮助。这是一个 Python 解决方案：

from collections import defaultdict
result = defaultdict(int)
for p,f in zip(positions,forces):
    result[p] += f

positions, forces = zip(*result.items())
print positions, forces

编辑： 我不确定 "I have to do it with numpy" 是什么意思，但是

import numpy as np
positions = np.array([1,1,4,4])
forces = np.array([4,5,8,9])
up = np.unique(positions)
uf = np.fromiter((forces[positions == val].sum() for val in up), dtype=int)

print up, uf

使用 NumPy 对数组中的元素求和

Sum elements in array with NumPy

python

arrays

numpy

matplotlib

编辑