A question about nested indexing of `numpy` arrays

Question

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.

One way of thinking of this operation is that you "replaced" each element of arr2 with the corresponding row of `arr1'. Hence the (4,8) shape has expanded to (4,8,5). — hpaulj
– hpaulj, Commented Jan 11, 2023 at 17:11

Chrysophylaxs · Accepted Answer · 2023-01-11 11:47:20Z

1

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.

answered Jan 11, 2023 at 11:47

Chrysophylaxs

6,5933 gold badges13 silver badges25 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

James Arten Over a year ago

much clearer now, but why the shape has an extra dimension?

Chrysophylaxs Over a year ago

Because arr2 has shape (4, 8), so it defines a subspace with that same shape. The third dimension is resulting from the slice that is used to index the other axis in arr1 (with shape (5,). Together the result is shape (4, 8, 5).

Collectives™ on Stack Overflow

A question about nested indexing of `numpy` arrays

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related