Suppose I have a matrix A
of order m
×n
and a vector of order m
×1
. I would like to extract elements from each row of the matrix A
by using the elements of the vector as an offset in each row.
For example,
A = [[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]]
and a vector
y = [4, 2, 1, 3, 2]
What I want to achieve is a way to extract the elements of A
such that each element of the vector indexes an element in the corresponding row of A
, i.e., implementing
for i in range(len(y)):
A[i, y[i]] = #perform operations here
without the use of any explicit loops.
The expected output is,
[3, 2, 5, 1, 1]
I am using Python and the NumPy library.
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]])