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:
In [270]: import numpy as np
In [271]: 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]])
In [272]: cols = [4, 2, 1, 3, 2]
And after that, nothing prevents you from using advanced indexing:
In [273]: rows = np.arange(A.shape[0])
In [274]: rows
Out[274]: array([0, 1, 2, 3, 4])
In [275]: A[rows, cols]
Out[275]: array([3, 2, 5, 1, 1])
In [276]: A[rows, cols] = -99
In [277]: A
Out[277]:
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]])