Самый эффективный способ уменьшить-суммировать массив numpy (с помощью autograd)

Question

У меня есть два массива:

index  = [2,1,0,0,1,1,1,2]
values = [1,2,3,4,5,4,3,2]

Я хотел бы произвести:

[sum(v for i,v in zip(index, values) if i == ui) for i in sorted(set(index))]

наиболее эффективным способом.

• мои значения вычисляются через автоград
• делать групповой в пандах действительно неэффективно из-за пункта выше
• Я должен сделать это сотни раз на одном и том же index но с разными значениями
• len(values) ~ 10 ** 7
• len(set(index)) ~ 10 ** 6
• Counter(index).most_common(1)[0][1] ~ 1000

Я думаю, что чистое решение было бы лучшим.

Я пытался предварительно вычислить сокращенную версию indexи затем выполните:

[values[l].sum() for l in reduced_index]

но это не достаточно эффективно.

Вот минимальный пример кода:

import numpy as np
import autograd.numpy as anp
from autograd import grad
import pandas as pd

EASY = True

if EASY:
    index = np.random.randint(10, size=10**3)
    values = anp.random.rand(10**3) * 2 - 1
else:
    index = np.random.randint(1000, size=10**7)
    values = anp.random.rand(10**7) * 2 - 1


# doesn't work
def f1(values):
    return anp.exp(anp.bincount(index, weights=values)).sum()


index_unique = sorted(set(index))
index_map = {j: i for i, j in enumerate(index_unique)}
index_mapped = [index_map[i] for i in index]
index_lists = [[] for _ in range(len(index_unique))]
for i, j in enumerate(index_mapped):
    index_lists[j].append(i)


def f2(values):
    s = anp.array([values[l].sum() for l in index_lists])
    return anp.exp(s).sum()


ans = grad(f2)(values)

Answer 1

Если ваш индекс не отрицательные целые числа, вы можете использовать np.bincount с values в качестве весов:

np.bincount(index, weights=values)
# array([ 7., 14.,  3.])

Это дает сумму в каждой позиции от 0 в max(index),