A question about nested indexing of `numpy` arrays

January 11, 2023

I’m trying to understand what the following does at a conceptual level. Let’s say we have two numpy arrays of random integers

arr1

array([[2, 2, 2, 2, 1],
       [1, 3, 1, 3, 2],
       [2, 2, 2, 1, 3],
       [1, 1, 1, 3, 2]])

arr2

array([[1, 3, 1, 1, 3, 3, 2, 2],
       [2, 3, 2, 2, 2, 3, 2, 1],
       [3, 3, 3, 1, 1, 3, 3, 3],
       [1, 1, 2, 1, 2, 1, 1, 1]])

Then, I do a nested indexing of the second array arr2 into the first one arr1, obtaining

arr1[arr2,:]

array([[[1, 3, 1, 3, 2],
        [1, 1, 1, 3, 2],
        [1, 3, 1, 3, 2],
        [1, 3, 1, 3, 2],
        [1, 1, 1, 3, 2],
        [1, 1, 1, 3, 2],
        [2, 2, 2, 1, 3],
        [2, 2, 2, 1, 3]],

       [[2, 2, 2, 1, 3],
        [1, 1, 1, 3, 2],
        [2, 2, 2, 1, 3],
        [2, 2, 2, 1, 3],
        [2, 2, 2, 1, 3],
        [1, 1, 1, 3, 2],
        [2, 2, 2, 1, 3],
        [1, 3, 1, 3, 2]],

       [[1, 1, 1, 3, 2],
        [1, 1, 1, 3, 2],
        [1, 1, 1, 3, 2],
        [1, 3, 1, 3, 2],
        [1, 3, 1, 3, 2],
        [1, 1, 1, 3, 2],
        [1, 1, 1, 3, 2],
        [1, 1, 1, 3, 2]],

       [[1, 3, 1, 3, 2],
        [1, 3, 1, 3, 2],
        [2, 2, 2, 1, 3],
        [1, 3, 1, 3, 2],
        [2, 2, 2, 1, 3],
        [1, 3, 1, 3, 2],
        [1, 3, 1, 3, 2],
        [1, 3, 1, 3, 2]]])

which is a new array with shape (4,8,5). It is not clear to me how should I interpret this new object, and how the entries of the two arrays are actually combined together.

>Solution :

Reference on numpy ndarray indexing with integer arrays

TLDR:

out = arr1[arr2, :]
out[i, j, k] == arr1[ arr2[i, j], k ] # for all valid indices i,j,k

Intuition:

The values inside arr2 are being used independently/separately to index the first axis of arr1, and the results are placed into a new array with the same shape as arr2.