How to extract elements from a matrix using a vector of indices?
You should start by converting list A
into a NumPy array:
>>> import numpy as np
>>> A = np.array([[3, 0, 0, 8, 3],
... [9, 3, 2, 2, 6],
... [5, 5, 4, 2, 8],
... [3, 8, 7, 1, 2],
... [3, 9, 1, 5, 5]])
...
>>> y = [4, 2, 1, 3, 2]
And after that, nothing prevents you from using advanced indexing:
>>> A[np.arange(A.shape[0]), y]
array([3, 2, 5, 1, 1])
>>> A[np.arange(A.shape[0]), y] = -99
>>> A
array([[ 3, 0, 0, 8, -99],
[ 9, 3, -99, 2, 6],
[ 5, -99, 4, 2, 8],
[ 3, 8, 7, -99, 2],
[ 3, 9, -99, 5, 5]])