Search code examples
pythonnumpymatrixmultidimensional-arraynumpy-indexing

Advanced Integer slicing when slicing object is an ndarray tuple


I understand how

x=np.array([[1, 2], [3, 4], [5, 6]] 
y = x[[0,1,2], [0,1,0]] 

Output gives y= [1 4 5] This just takes the first list as rows and seconds list and columns.

But how does the the below work?

x = np.array([[ 0,  1,  2],[ 3,  4,  5],[ 6,  7,  8],[ 9, 10, 11]]) 

rows = np.array([[0,0],[3,3]])
cols = np.array([[0,2],[0,2]]) 
y = x[rows,cols] 

This gives the output of :

[[ 0  2]                                                                      
 [ 9 11]] 

Can you please explain the logic when using ndarrays as slicing object? Why does it have a 2d array for both rows and columns. How are the rules different when the slicing object is a ndarray as opposed to a python list?


Solution

  • We've the following array x

    x = np.array([[1, 2], [3, 4], [5, 6]] 
    

    And the indices [0, 1, 2] and [0, 1, 0] which when indexed into x like

    x[[0,1,2], [0,1,0]] 
    

    gives

    [1,  4,  5]
    

    The indices that we used basically translates to:

    [0, 1, 2] & [0, 1, 0] --> [0,0], [1,1], [2,0]
    

    Since we used 1D list as indices, we get 1D array as result.


    With that knowledge, let's see the next case. Now, we've the array x as:

    x = np.array([[ 0,  1,  2],[ 3,  4,  5],[ 6,  7,  8],[ 9, 10, 11]]) 
    

    Now the indices are 2D arrays.

    rows = np.array([[0,0],[3,3]])
    cols = np.array([[0,2],[0,2]]) 
    

    This when indexed into the array x like:

    x[rows,cols] 
    

    simply translates to:

    [[0,0],[3,3]]
      | |   | |      ====> [[0,0]], [[0,2]], [[3,0]], [[3,2]]
    [[0,2],[0,2]]
    

    Now, it's easy to observe how these 4 list of list when indexed into the array x would give the following result (i.e. here it simply returns the corner elements from our array x):

    [[ 0,  2]                                                                      
     [ 9, 11]]
    

    Note that in this case we get the result as a 2D array (as opposed to 1D array in the first case) since our indices rows & columns itself were 2D arrays (i.e. equivalently list of list) whereas in the first case our indices were 1D arrays (or equivalently simple list without any nesting).

    So, if you need 2D arrays as result, you need to give 2D arrays as indices.