根据偏移量更改numpy数组的值

我有一个二维的numpy数组

对于每一列,我想在第N行中加1,其中N是该列第0行中的值。

我怎样才能做到这一点?

A=np.zeros(100)
A=np.reshape(A,[20,5])
A[0]=[5,2,4,1,3]

I want to add 1 to A[5,0], A[2,1], A[4,2], A[1,3] and A[3,4]

评论
  • ab_non
    ab_non 回复

    Following the following logic: add 1 to A[5,0], A[2,1], A[4,2], A[1,3] and A[3,4], you can do so using advanced indexing:

    indices = [5,2,4,1,3]
    A[indices, np.arange(len(indices))] = 1
    
    print(A)
    array([[0., 0., 0., 0., 0.],
           [0., 0., 0., 1., 0.],
           [0., 1., 0., 0., 0.],
           [0., 0., 0., 0., 1.],
           [0., 0., 1., 0., 0.],
           [1., 0., 0., 0., 0.],
           [0., 0., 0., 0., 0.],
           ...