NumPy - Vectorizing bincount over 2D array column wise with weights

Question

I've been looking at the solutions here and here but failing to see how I can apply it to my structures.

I have 3 arrays: an (M, N) of zeros, and (P,) of indexes (some repeat) and an (P, N) of values.

I can accomplish it with a loop:

# a: (M, N)
# b: (P, N)
# ix: (M,)
for i in range(N):
    a[:, i] += np.bincount(ix, weights=b[:, i], minlength=M)

I've not seen any examples that use indexes in this manner, or with the weights keyword. I understand I need to bring everything into a 1D array to vectorize it, however I am struggling to figure out how to accomplish that.

Answer 1

Basic idea stays the same as discussed in some detail in those linked posts, ie create a 2D array of bins with offsets per "1D data" to be processed (per col in this case). So, with those in mind, we will end up with something like this -

# Extent of bins per col
n = ix.max()+1

# 2D bins for per col processing
ix2D = ix[:,None] + n*np.arange(b.shape[1])

# Finally use bincount with those 2D bins as flattened and with
# flattened b as weights. Reshaping is needed to add back into "a".
a[:n] += np.bincount(ix2D.ravel(), weights=b.ravel(), minlength=n*N).reshape(N,-1).T