pythonnumpymatrixvectormatrix-indexing

How to extract elements from a matrix using a vector of indices?


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.


Solution

  • 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]])