I would like to use a numpy array to index a numpy matrix using the values of the array indicating the columns, and indices indicating the corresponding row numbers.As an example, I have a numpy matrix,

a = np.tile(np.arange(1920), (41, 1))
>>> [[0, 1, 2, ..., 1919]
     [0, 1, 2, ..., 1919]
     ... 
     [0, 1, 2, ..., 1920]]
b = np.arange(40, -1, -1) # We want to do a[b] in the most efficient way.

What I would like to get is an array c which is,

c = [40, 39, 38, ..., 0]

That is, I want to use b to get the following indices from a,

[(0, b[0]), (1, b[1]), ... (40, b[40])] # 0th row b[0]th column, 1st row b[1]th column...

How do I do this, and what is the most efficient way to do this?

>Solution :

You can use advanced indexing with 2 integer arrays. For the 0th axis ("the rows"), you simply use 0..n (can be generated using np.arange) with n the length of b. For the 1st axis ("the columns"), you use b:

import numpy as np

# Setup:
a = np.tile(np.arange(1920), (41, 1))
b = np.arange(40, -1, -1)

# Solution:
c = a[np.arange(len(b)), b]

c:

array([40, 39, 38, 37, 36, 35, 34, 33, 32, 31, 30, 29, 28, 27, 26, 25, 24,
       23, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10,  9,  8,  7,
        6,  5,  4,  3,  2,  1,  0])