pythonpython-3.xnumpynumpy-indexingadvanced-indexing

How this numpy advance indexing code works?


I am learning numpy framework.This piece of code I don't understand.

import numpy as np
a =np.array([[0,1,2],[3,4,5],[6,7,8],[9,10,11]])
print(a)
row = np.array([[0,0],[3,3]])
col = np.array([[0,2],[0,2]])
b = a[row,col]
print("This is b array:",b)

This b array returns the corner values of a array, that is, b equals [[0,2],[9,11]].


Solution

  • You're indexing a using two equally shaped 2d-arrays, hence you're output array will also have the same shape as col and row. To better understand how array indexing works you can check the docs, where as shown, indexing with 1d-arrays over the existing axis' of a given array works as follows:

    result[i_1, ..., i_M] == x[ind_1[i_1, ..., i_M], ind_2[i_1, ..., i_M], ..., ind_N[i_1, ..., i_M]]

    Where the same logic applies in the case of indexing with 2d-arrays over each axis, but instead you'd have a result array with up to i_N_M indices.

    So going back to your example you are essentially selecting from the rows of a based on row, and from those rows you are selecting some columns col. You might find it more intuitive to translate the row and column indices into (x,y) coordinates:

    (0,0), (0,2)   
    (3,0), (3,2)   
    

    Which, by accordingly selecting from a, results in the output array:

    print(a[row,col])
    
    array([[ 0,  2],
           [ 9, 11]])