Numpy: How can I index the rows of one matrix using indices found in the rows of another?

Question

I have two (2000, 10) matrices: weight_values contains a set of values and weight_indexes contains a set of integers to be used as indexes to a new matrix.

I would like to use weight_indexes to select entries from a new (2000, 2000) zero matrix and then set those columns to be the values found in the value matrix.

For example, doing this gets me what I want:

weights = np.zeros((2000, 2000))
    
for i in range(weight_indexes.shape[0]):
    weights[i, weight_indexes[i]] = weight_values[i]

However when I try doing this using array indexing it doesn't work. Indexing weights using weight_indexes like this:

weights[:, weight_indexes[:]]

...rather than selecting the appropriate columns from weights , this creates a new (2000, 2000, 10) sized matrix.

Is there some vectorised way I can do this without using loops?

Answer 1

Python fancy indices broadcast together. If you want to set 10 elements in each row of weights , make i broadcast to the same shape as weight_indexes :

weights[np.arange(len(weight_indexes))[:, None], weight_indexes] = weight_values