Сгруппировать Python по массиву a и суммировать массив b - Производительность

Question

Учитывая два неупорядоченных массива одинаковой длины a и b:

a = [7,3,5,7,5,7]
b = [0.2,0.1,0.3,0.1,0.1,0.2]

Я бы хотел сгруппировать по элементам в a:

aResult = [7,3,5]

, суммируя по элементам в b(Пример, использованный для суммирования функции плотности вероятности):

bResult = [0.2 + 0.1 + 0.2, 0.1, 0.3 + 0.1] = [0.5, 0.1, 0.4]

Альтернативно, случайные a и b в python:

import numpy as np
a = np.random.randint(1,10,10000)
b = np.array([1./len(a)]*len(a))

У меня есть два подхода, которые наверняка далеки отнижняя граница производительности.Подход 1 (по крайней мере, хороший и короткий): Время: 0,769315958023

def approach_2(a,b):
    bResult = [sum(b[i == a]) for i in np.unique(a)]
    aResult = np.unique(a)

Подход 2 (numpy.groupby, ужасно медленный) Время: 4.65299129486

def approach_2(a,b): 
    tmp = [(a[i],b[i]) for i in range(len(a))]
    tmp2 = np.array(tmp, dtype = [('a', float),('b', float)])
    tmp2 = np.sort(tmp2, order='a') 

    bResult = []
    aResult = []
    for key, group in groupby(tmp2, lambda x: x[0]):
        aResult.append(key)
        bResult.append(sum([i[1] for i in group]))

Обновление: подход 3, Пабло,Время: 1.0265750885

def approach_Pablo(a,b):    

    pdf = defaultdict(int); 
    for x,y in zip(a,b):
        pdf[x] += y  

Обновление: Подход 4, Unutbu.Время: 0.184849023819 [ПОБЕДИТЕЛЬ ТАК ФАР, но только как целое число]

def unique_Unutbu(a,b):

    x=np.bincount(a,weights=b)
    aResult = np.unique(a)
    bResult = x[aResult]

Может быть, кто-то найдет более разумное решение этой проблемы, чем я:)

Answer 1

Этот подход похож на @ unutbu's one :

import numpy as np

def f(a, b):
    result_a, inv_ndx = np.unique(a, return_inverse=True)
    result_b = np.bincount(inv_ndx, weights=b)
    return result_a, result_b

Он допускает нецелочисленный тип для массива a.Это позволяет большие значения в массиве a.Возвращает a элементов в отсортированном порядке.Если желательно, легко восстановить исходный порядок, используя аргумент return_index функции np.unique().

Производительность ухудшается с увеличением количества уникальных элементов в a.Это в 4 раза медленнее, чем версия @ unutbu по данным вашего вопроса.

Я сделал сравнение производительности с тремя дополнительными методами.Лидерами являются: для целочисленных массивов - реализация на основе хеша в Cython;для double массивов (для входного размера 10000) - на основе сортировки импл. также в Cython.

Answer 2

Если a состоит из целых чисел <2 ** 31-1 (то есть, если <code>a имеет значения, которые могут вписаться в dtype int32), тогда вы можете использовать np.bincount с весами:

import numpy as np
a = [7,3,5,7,5,7]
b = [0.2,0.1,0.3,0.1,0.1,0.2]

x=np.bincount(a,weights=b)
print(x)
# [ 0.   0.   0.   0.1  0.   0.4  0.   0.5]

print(x[[7,3,5]])
# [ 0.5  0.1  0.4]

np.unique(a) возвращает [3 5 7], поэтому результат отображается в другом порядке:

print(x[np.unique(a)])
# [ 0.1  0.4  0.5]

Одна потенциальная проблема при использовании np.bincount состоит в том, что он возвращает массив, длина которого равнаравно максимальному значению в a.Если a содержит хотя бы один элемент со значением около 2 ** 31-1, то bincount должен будет выделить массив размером 8*(2**31-1) байт (или 16 ГБ).

Так что np.bincount можетбыть самым быстрым решением для массивов a, которые имеют большую длину, но не большие значения.Для массивов a, которые имеют небольшую длину (и большие или малые значения), использование collections.defaultdict, вероятно, будет более быстрым.

Редактирование: См. решение Дж. Ф. Себастьяна для обходаограничение только целых значений и проблема больших значений.

Answer 3

Как насчет этого подхода:

from collections import defaultdict
pdf = defaultdict(int)
a = [7,3,5,7,5,7]
b = [0.2,0.1,0.3,0.1,0.1,0.2]
for x,y in zip(a,b):
  pdf[x] += y

Вы просматриваете каждый элемент только один раз и используете словарь для быстрого поиска.Если вам действительно нужны два отдельных массива, вы можете запросить их:

aResult = pdf.keys()
bResult = pdf.values()