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.我通过{获得它们{1}}。

示例的预期输出应为：

x[n[0][0]],x[n[0][1]] and x[n[0][2]]

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

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

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

由于这是迭代的，我遍历列表x和n。以下代码实现了这一点，但在列表x和n中逐个元素地进行。

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

我想让这段代码更具矢量性，但这是几天前我最好的。

编辑：如果它有用，我总是每个元素总计3个值，因此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。

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)

Answer 2

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

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

如何改进这个python矢量函数

2 个答案: