python-将bincount应用于2D numpy数组的每一行

有没有办法使用“轴= 1”应用bincount?所需的结果将与列表理解相同:

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]])

最佳答案

np.bincount不适用于沿着特定轴的2D阵列.为了通过对np.bincount的单个矢量化调用获得理想的效果,可以创建一维ID数组,以便即使元素相同,不同的行也将具有不同的ID.当使用一次具有这些ID的np.bincount调用时,这将使不同行中的元素不会合并在一起.因此,可以在考虑线性索引的思想的情况下创建此类ID数组,如下所示-

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

然后,将ID输入np.bincount并最终重新调整为2D-

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