Vectorize добавление в массив индексируется другой массив

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] добавляется только один раз, не дважды). Кроме того, это, похоже, не быстрее, чем цикл без использования.

Есть ли лучший способ сделать это?

EDIT: YXD связанный this answer в комментариях, которые поначалу, похоже, соответствуют счету. Казалось бы, вы можете делать то, что я хочу с

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

и он работает, но когда я пытаюсь запустить его это значительно медленнее, чем unvectorized версии. Итак, остается вопрос: есть ли более эффективный способ сделать это?

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) 

для цикла, при запуске с помощью %timeit в IPython дает на моей машине:

10 loops, best of 3: 19.4 ms per loop 

версия subtract.at производит такое же значение для A в конце, но намного медленнее:

1 loops, best of 3: 444 ms per loop 

Answer 1

Код для исходного для цикла на основе подхода будет выглядеть примерно так -

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

---

Case # 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, модифицированный для цикла подход с меньшей сложностью контура может быть предложено -

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