Performing operation on 2D array using indices from 1D array-CodePudding

I have the following array in python:

a = np.array([[1,1,1],[1,1,1],[1,1,1]])

and the following index array:

b = np.array([0,1,2])

I want to index a using b such that I can subtract 1 from the matching row/column and get the following result:

[[0,1,1],[1,0,1],[1,1,0]]

I can do it using loops, wanted to know if there was a "non-loop" way of doing it.

for i in range(len(b)): 
    a[i][b[i]] = a[i][b[i]] - 1

CodePudding user response：

Your code produces output as follows:

a = np.array([[1,1,1],[1,1,1],[1,1,1]])
b = np.array([0,1,2])
for i in range(len(b)): 
   a[i][b[i]] = a[i][b[i]] - 1

Output:

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

This can be done in non -loopy way as follows:

a[np.arange(len(a)),b] -= 1
print(a)

Output:

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

CodePudding user response：

It looks like there is some confusion on how to handle this.

You want a simple indexing:

a[np.arange(len(a)), b] -= 1

Output:

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

Output for b = np.array([2,0,1])

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