python - 将 bincount 应用于二维 numpy 数组的每一行

Question

有没有办法将 bincount 与“axis = 1”一起应用？期望的结果将与列表理解相同:

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不适用于沿特定轴的二维数组。要通过对 np.bincount 的单个矢量化调用获得所需的效果，可以创建一个一维 ID 数组，这样即使元素相同，不同的行也会有不同的 ID。当使用这些 ID 对 np.bincount 进行一次调用时，这将使来自不同行的元素不会合并在一起。因此，可以根据 linear indexing 的想法创建这样的 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)