使用NumPy

时间:2015-05-06 17:25:05

标签: python arrays numpy matplotlib

我用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处产生两个箭头,在彼此之上。相反,我想要对力进行求和,并且仅在力4 + 5 = 9时在位置1处创建一个箭头,在力4 + 9 = 17时在位置4处创建一个箭头。

如何使用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

1 个答案:

答案 0 :(得分:3)

我不确定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

修改 我不确定“我必须用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