如何改进这个python向量函数

Question

我有一个函数A，它的输入是numpy向量(numpy.ndarray公司)称为x。此函数为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]] and 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

您至少可以按如下方式移除内环：

def A(x, n):
    a = 3 * x
    for i in range(len(x)):
        a[i] += np.sum(x[np.ix_(n[i]-1)], axis=0)
    return a

Answer 2

（有关更简单、更快速的解决方案，请参见末尾的更新）

这可以通过广播技术在没有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行代码简化为一行代码（检查下面的注释），速度也提高了20%。你知道吗

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