python - 如何改进这个python矢量函数

Question

我有一个函数 A，它的输入是一个名为 x 的 numpy 向量 (numpy.ndarray)。此函数为 x 的每个元素计算该元素本身与 x 的其他元素的总和，这些元素由这些元素的列表给出。

下面的例子应该能更好地说明这一点:

x = [[2,3], [3,4], [1,2], [1,3], [1,4]] # my input
n = [[1,2,3], [0,4,2], [3,0,1], [0,1,4], [3,1,2]] # list with lists of element to be added for each element in x

所以对于 x 的第一个元素，即 x[0] = [2,3]，我必须添加 n[0] 给定的值，所以它们是 1、2 和 3。我通过 x[n[0][0]]、x[n[0][1]] 和 x[n[0][2]]。

该示例的预期输出应该是:

l = [[11, 18], [13, 21], [9, 16], [9, 20], [8, 21]]

元素 x[i] 的最终总和应该是

(x[i] + x[n[i][0]] + x[i] + x[n[i][1]] + x[i] + x[n[i][2]])

函数的返回值是每个计算总和的列表。

因为这是迭代的，所以我遍历了列表 x 和 n。以下代码实现了这一点，但在列表 x 和 n 中逐个元素。

def A(x):
    a = []
    for i, x_i in enumerate(x):
        mysum = np.zeros(2)
        for j, n_j in enumerate(n[i]):
           mysum = mysum + x_i + x[n_j]
        a.append(mysum)
    return np.array(a)

我想让这段代码更加矢量化，但这是几天前以来我最好的。

编辑:如果有帮助，我总是对每个元素求和 3 个值，因此 n 的子列表的长度始终为 3。

Answer 1

(请参阅末尾的更新以获得更简单、更快速的解决方案)

这可以通过广播技术在没有 for 循环的情况下完成

def C(x,n):
    y = x[n.ravel()-1]
    z = y.reshape((-1,3,2))
    xx = x[:,np.newaxis,:]
    ans = z+xx
    ans = ans.sum(axis=1)
    return ans

与使用 for 循环的解决方案相比，它至少快 5-6 倍。

In [98]: np.all(A(x,n)==C(x,n))
Out[98]: True

In [95]: %timeit ans=A(x,n)
10000 loops, best of 3: 153 us per loop

In [96]: %timeit ans=C(x,n)
10000 loops, best of 3: 27 us per loop

更新

Jaime 将我的 6 行代码减少为简单的 1 行代码(查看下面的评论)，而且速度也提高了 20%。

ans = 3*x + x[n-1].sum(axis=1)