Python group by array a, and summarize array b - Performance

Question

Given two unordered arrays of same lengths a and b: ``` a = [7,3,5,7,5,7] b = [0.2,0.1,0.3,0.1,0.1,0.2] ``` I'd like to group by the elements in a: ``` aResult = [7,3,5] ``` summing over the elements in b (Example used to summarize a probability density function): ``` bResult = [0.2 + 0.1 + 0.2, 0.1, 0.3 + 0.1] = [0.5, 0.1, 0.4] ``` Alternatively, random a and b in python: ``` import numpy as np a = np.random.randint(1,10,10000) b = np.array([1./len(a)]*len(a)) ``` I have two approaches, which for sure are far from the lower performance boundary. Approach 1 (at least nice and short): Time: 0.769315958023 ``` def approach_2(a,b): bResult = [sum(b[i == a]) for i in np.unique(a)] aResult = np.unique(a) ``` Approach 2 (numpy.groupby, horribly slow) Time: 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…

Accepted Answer

If `a` is composed of ints < 2**31-1 (that is, if `a` has values that can fit in dtype `int32`), then you could use `np.bincount` with weights:

```
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)` returns `[3 5 7]`, so the result appears in a different order:

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

One potential problem with using `np.bincount` is that it returns an array whose length is equal to the maximum value in `a`. If `a` contains even one element with value near 2**31-1, then `bincount` would have to allocate an array of size `8*(2**31-1)` bytes (or 16GiB).

So `np.bincount` might be the fastest solution for arrays `a` which have big length, but not big values. For arrays `a` which have small length (and big or small values), using a `collections.defaultdict` would probably be faster.

Edit: See J.F. Sebastian's solution for a way around the integer-values-only restriction and big-values problem.

Python group by array a, and summarize array b - Performance

Solution

Problem