>>> import numpy_indexed as npi
>>> import numpy as np
>>> a = np.array([[0,0,1,1,2,2], [4,4,8,8,10,10]]).T
>>> a
array([[ 0, 4],
[ 0, 4],
[ 1, 8],
[ 1, 8],
[ 2, 10],
[ 2, 10]])
>>> npi.group_by(a[:, 0]).sum(a[:,1])
(array([0, 1, 2]), array([ 8, 16, 20], dtype=int32))
I want to perform calculations on subsets of the second column clustered by the first column on large sets (~1m lines). Is there an efficient (and/or vectorised) way to use the output of group_by by numpy_indexed in order to add a new column with the output of these calculations? In the example of sum as above I would like to produce the output below.
If there is an efficient way of doing this without using numpy_indexed in the first place, that would also be very helpful.
array([[ 0, 4, 8],
[ 0, 4, 8],
[ 1, 8, 16],
[ 1, 8, 16],
[ 2, 10, 20],
[ 2, 10, 20]])