NumPy: Выбрать и суммировать данные в массив

Question

У меня есть (большой) массив данных и (большой) список списков (несколько) индексов, например,

data = [1.0, 10.0, 100.0] 
contribs = [[1, 2], [0], [0, 1]] 

Для каждой записи в contribs , Я хотел бы подытожить соответствующие значения data и поместить их в массив. Для приведенного выше примера, ожидаемый результат будет

out = [110.0, 1.0, 11.0] 

Делая это в цикле работ,

c = numpy.zeros(len(contribs)) 
for k, indices in enumerate(contribs): 
    for idx in indices: 
     c[k] += data[idx] 

, но так как data и contribs большие, это занимает слишком много времени.

У меня такое ощущение, что это можно улучшить, используя причудливую индексацию numpy.

Любые подсказки?

Answer 1

можно было бы,

data = np.array(data) 
out = [np.sum(data[c]) for c in contribs] 

Должен быть быстрее, чем двойная петля, по крайней мере.

Answer 2

Вот почти векторную * подход -

# Get lengths of list element in contribs and the cumulative lengths 
# to be used for creating an ID array later on. 
clens = np.cumsum([len(item) for item in contribs]) 

# Setup ID array that corresponds to same ID for same list element in contribs. 
# These IDs would be used to accumulate values from a corresponnding array 
# that is created by indexing into data array with a flattened contribs 
id_arr = np.zeros(clens[-1],dtype=int) 
id_arr[clens[:-1]] = 1 
out = np.bincount(id_arr.cumsum(),np.take(data,np.concatenate(contribs))) 

Этот подход включает в себя некоторые настройки на работу. Таким образом, преимущества можно было бы надеяться увидеть при подаче с приличными массивами ввода и приличным количеством элементов списка в contribs, что соответствовало бы циклу в обратном решении.

* Обратите внимание, что это придумано как , почти векторизованное, потому что только цикл, выполняемый здесь, находится в начале, где мы получаем длины элементов списка. Но эта часть, не требующая вычислительных требований, должна иметь минимальный эффект на общую продолжительность работы.

Answer 3

Я не уверен, что все случаи работают, но для примера, с data как numpy.array:

# Flatten "contribs" 
f = [j for i in contribs for j in i] 

# Get the "ranges" of data[f] that will be summed in the next step 
i = [0] + numpy.cumsum([len(i) for i in contribs]).tolist()[:-1] 

# Take the required sums 
numpy.add.reduceat(data[f], i)