如何改进这个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

(请参阅最后的更新,以获得更简单,更快速的解决方案)

这可以通过广播技术在没有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

UPDATE

Jaime已将我的6行代码缩减为简单的1行代码(请查看下面的评论),并且速度提高了20%.

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