python - 将加法矢量化到由另一个数组索引的数组中

Question

我正在尝试获得以下循环的快速矢量化版本:

for i in xrange(N1):
   A[y[i]] -= B[i,:]

这里A.shape = (N2,N3)，y.shape = (N1)，其中y取中的值[0，N2[，B.shape = (N1，N3)。您可以将 y 条目视为 A 行的索引。这里 N1 很大，N2 非常小，N3 很小。

我以为只是做

A[y] -= B

可以，但问题是 y 中有重复的条目，这并不能做正确的事情(即，如果 y=[1,1]那么 A[1] 仅添加一次，而不是两次)。而且这似乎并不比非矢量化的 for 循环快。

有更好的方法吗？

编辑:YXD 链接 this answer乍一看似乎符合要求的评论。看起来你完全可以做我想要的事情

np.subtract.at(A, y, B)

它确实有效，但是当我尝试运行它时，它比非矢量化版本明显慢。所以，问题仍然是:是否有更高效的方法来做到这一点？

EDIT2:一个例子，让事情变得具体:

n1,n2,n3 = 10000, 10, 500
A = np.random.rand(n2,n3)
y = np.random.randint(n2, size=n1)
B = np.random.rand(n1,n3)

当在 ipython 中使用 %timeit 运行时，for 循环在我的机器上给出:

10 loops, best of 3: 19.4 ms per loop

subtract.at 版本最终为 A 生成相同的值，但速度慢得多:

  1 loops, best of 3: 444 ms per loop

Answer 1

基于 for 循环的原始方法的代码看起来像这样 -

def for_loop(A):
    N1 = B.shape[0]
    for i in xrange(N1):
       A[y[i]] -= B[i,:]
    return A

---

案例#1

如果 n2 >> n3，我建议采用这种矢量化方法 -

def bincount_vectorized(A):

    n3 = A.shape[1]
    nrows = y.max()+1
    id = y[:,None] + nrows*np.arange(n3)
    A[:nrows] -= np.bincount(id.ravel(),B.ravel()).reshape(n3,nrows).T
    return A

运行时测试 -

In [203]: n1,n2,n3 = 10000, 500, 10
     ...: A = np.random.rand(n2,n3)
     ...: y = np.random.randint(n2, size=n1)
     ...: B = np.random.rand(n1,n3)
     ...: 
     ...: # Make copies
     ...: Acopy1 = A.copy()
     ...: Acopy2 = A.copy()
     ...: 

In [204]: %timeit for_loop(Acopy1)
10 loops, best of 3: 19 ms per loop

In [205]: %timeit bincount_vectorized(Acopy2)
1000 loops, best of 3: 779 µs per loop

---

案例#2

如果 n2 << n3，则可以建议采用循环复杂度较低的修改版 for 循环方法 -

def for_loop_v2(A):
    n2 = A.shape[0]
    for i in range(n2):
        A[i] -= np.einsum('ij->j',B[y==i]) # OR (B[y==i]).sum(0)
    return A

运行时测试 -

In [206]: n1,n2,n3 = 10000, 10, 500
     ...: A = np.random.rand(n2,n3)
     ...: y = np.random.randint(n2, size=n1)
     ...: B = np.random.rand(n1,n3)
     ...: 
     ...: # Make copies
     ...: Acopy1 = A.copy()
     ...: Acopy2 = A.copy()
     ...: 

In [207]: %timeit for_loop(Acopy1)
10 loops, best of 3: 24.2 ms per loop

In [208]: %timeit for_loop_v2(Acopy2)
10 loops, best of 3: 20.3 ms per loop