使用NumPy数组作为NumPy数组的索引

Question

I have a 3x3x3 NumPy array:

>>> x = np.arange(27).reshape((3, 3, 3))
>>> x
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])

Now I create an ordinary list of indices:

>>> i = [[0, 1, 2, 1], [2, 1, 0, 1], [1, 2, 0, 1]]

As expected, I get four values using this list as the index:

>>> x[i]
array([ 7, 14, 18, 13])

But if I now convert i into a NumPy array, I won't get the same answer.

>>> j = np.asarray(i)
>>> x[j]
array([[[[ 0,  1,  2],
         [ 3,  4,  5],
         [ 6,  7,  8]],

        [[ 9, 10, 11],
         [12, 13, 14],
         [15, 16, 17]],

        [[18, 19, 20],
         [21, 22, 23],
         [24, 25, 26]],

       ...,
       [[[ 9, 10, 11],
         [12, 13, 14],
         [15, 16, 17]],

        [[18, 19, 20],
         [21, 22, 23],
         [24, 25, 26]],

        [[ 0,  1,  2],
         [ 3,  4,  5],
         [ 6,  7,  8]],

        [[ 9, 10, 11],
         [12, 13, 14],
         [15, 16, 17]]]])

Why is this so? Why can't I use NumPy arrays as indices to NumPy array?

Answer 1

x[j] is the equivalent of x[j,:,:]

In [163]: j.shape
Out[163]: (3, 4)

In [164]: x[j].shape
Out[164]: (3, 4, 3, 3)

The resulting shape is the shape of j joined with the last 2 dimensions of x . j just selects from the 1st dimension of x .

x[i] on the other hand, is the equivalent to x[tuple(i)] , that is:

In [168]: x[[0, 1, 2, 1], [2, 1, 0, 1], [1, 2, 0, 1]]
Out[168]: array([ 7, 14, 18, 13])

In fact x(tuple(j)] produces the same 4 item array.

The different ways of indexing numpy arrays can be confusing.

Another example of how the shape of the index array or lists affects the output:

In [170]: x[[[0, 1], [2, 1]], [[2, 1], [0, 1]], [[1, 2], [0, 1]]]
Out[170]: 
array([[ 7, 14],
       [18, 13]])

Same items, but in a 2d array.

Answer 2

Check out the docs for numpy , what you are doing is "Integer Array Indexing", you need to pass each coordinate in as a separate array:

j = [np.array(x) for x in i]

x[j]
Out[191]: array([ 7, 14, 18, 13])