Apply bincount to each row of a 2D numpy array

Question

Is there a way to apply bincount with "axis = 1"? The desired result would be the same as the list comprehension:

import numpy as np
A = np.array([[1,0],[0,0]])
np.array([np.bincount(r,minlength = np.max(A) + 1) for r in A])

#array([[1,1]
#       [2,0]])

Answer 1

np.bincount doesn't work with a 2D array along a certain axis. To get the desired effect with a single vectorized call to np.bincount, one can create a 1D array of IDs such that different rows would have different IDs even if the elements are the same. This would keep elements from different rows not binning together when using a single call to np.bincount with those IDs. Thus, such an ID array could be created with an idea of linear indexing in mind, like so -

N = A.max()+1
id = A + (N*np.arange(A.shape[0]))[:,None]

Then, feed the IDs to np.bincount and finally reshape back to 2D -

np.bincount(id.ravel(),minlength=N*A.shape[0]).reshape(-1,N)