Numpy накапливает один массив в другом, используя индексный массив

Мой вопрос касается конкретной операции с массивом, которую я хочу выразить с помощью numpy.

У меня есть массив floats w и массив индексов idx той же длины, что и w, и я хочу подвести все w с тем же значением idx и собрать их в массив v. Как цикл, это выглядит так:

for i, x in enumerate(w):
     v[idx[i]] += x

Есть ли способ сделать это с помощью операций с массивами? Моя догадка была v[idx] += w, но это не работает, поскольку idx содержит один и тот же индекс несколько раз.

Спасибо!

Ответ 1

numpy.bincount был введен для этой цели:

tmp = np.bincount(idx, w)
v[:len(tmp)] += tmp

Я думаю, что с 1.6 вы также можете передать minlength в bincount.

Ответ 2

Это известное поведение и, хотя и несколько неудачно, не имеет обходного режима на уровне numpy. (bincount можно использовать для этого, если вы крутите руку). Выполнение цикла самостоятельно - это ваш лучший выбор.

Обратите внимание, что ваш код мог быть более понятным без повторного использования имени w и без введения другого набора индексов, например

for i, w_thing in zip(idx, w):
    v[i] += w_thing

Если вам нужно ускорить этот цикл, вам может потребоваться спуститься до C. Cython делает это относительно легко.