Cfis Yoi Cfis Yoi - 1 month ago 9
Python Question

Python Update a submatrix given the indexes by a 2d matrix

Simple enough implementation with a for loop. Example:

s = np.zeros((3,3))
b = np.array(([0,1],[0,1],[2,0]))
for line in range(s.shape[0]):
s[line][b[line]] = 1


What i would like to do is to update the matrix without the for loop. Something like

s[b] = 1


But the s[b].shape returns a (3,2,3) matrix.
Is there any way to fix that or i have to make a for loop for every line?

Answer
s[np.arange(s.shape[0])[..., None], b] = 1

We use some tricky advanced indexing here. The np.arange(s.shape[0])[..., None] produces the column vector

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

which broadcasts against b to select, for each row of s, the column indices given by the elements of that row of b.

Comments