Increment given indices in a matrix

后端 未结 2 1123
不知归路
不知归路 2020-12-03 23:40

Briefly: there is a similar question and the best answer suggests using numpy.bincount. I need the same thing, but for a matrix.

I\'ve got two arrays:

2条回答
  •  时光取名叫无心
    2020-12-04 00:18

    You can still use bincount(). The trick is to convert a and b into a single 1D array of flat indices.

    If the matrix is nxm, you could apply bincount() to a * m + b, and construct the matrix from the result.

    To take the example in your question:

    In [15]: a = np.array([1, 2, 1, 1, 2])
    
    In [16]: b = np.array([2, 1, 1, 1, 1])
    
    In [17]: cnt = np.bincount(a * 3 + b)
    
    In [18]: cnt.resize((3, 3))
    
    In [19]: cnt
    Out[19]: 
    array([[0, 0, 0],
           [0, 2, 1],
           [0, 2, 0]])
    

    If the shape of the array is more complicated, it might be easier to use np.ravel_multi_index() instead of computing flat indices by hand:

    In [20]: cnt = np.bincount(np.ravel_multi_index(np.vstack((a, b)), (3, 3)))
    
    In [21]: np.resize(cnt, (3, 3))
    Out[21]: 
    array([[0, 0, 0],
           [0, 2, 1],
           [0, 2, 0]])
    

    (Hat tip @Jaime for pointing out ravel_multi_index.)

提交回复
热议问题