numpy alternative to successive indexing

Dom Byrne 0
Viewed 51

I'm sure this is probably a duplicate, but I couldn't find another question with exactly what I wanted so apologies in advance.

Given the following example array:

array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]],

       [[12, 13],
        [14, 15],
        [16, 17]]])

I want to find an equivalent of:

idx0, idx1 = [0,1], [0,2]
arr[idx0,...][:,idx1,:]

Which returns:

array([[[ 0,  1],
        [ 4,  5]],

       [[ 6,  7],
        [10, 11]]])

I'm sure there's a better/more concise way of doing this, but I can't find it.

indexing numpy arrays python
1 Answers

Is this what you're after?

import numpy as np

a = np.array(
    [[[ 0,  1],
      [ 2,  3],
      [ 4,  5]],

     [[ 6,  7],
      [ 8,  9],
      [10, 11]],

     [[12, 13],
      [14, 15],
      [16, 17]]]
)

result = a[0:2, [0, 2]]
print(result)

Result:

[[[ 0  1]
  [ 4  5]]

 [[ 6  7]
  [10 11]]]

The 0:2 part of the indexing selects the first two elements at the top level, the [0, 2] part selects only the first and third row of the results at the next level.

In the comments you clarified and asked about the specific difference between these two approaches:

result1 = a[0:2, [0, 2]]
result2 = a[[0, 1], [0, 2]]

It's a good question, because at face value, [0, 1] comes down to the same result as [0:2] in this case, after all:

print(a[0:2] == a[[0, 1]])              # prints an all `True` array
print(a[0:2].shape == a[[0, 1]].shape)  # prints `True`

However, the documentation states: "Note that in Python, x[(exp1, exp2, ..., expN)] is equivalent to x[exp1, exp2, ..., expN]; the latter is just syntactic sugar for the former."

And that's the issue here.

Compare:

result1 = a[[0, 1]][:, [0, 2]]  # works as needed
result2 = a[0:2, [0, 2]]        # same

However,

result3 = a[[0, 1], [0, 2]]     # this
result5 = a[[(0, 1), (0, 2)]]   # is just the same as this

So, you're just selecting the 0, 0 and 1, 2 elements.

To use the syntax from your example, for a correct result:

idx0, idx1 = slice(0, 2), [0, 2]
result = a[idx0, idx1]
print(result)

Result:

[[[ 0  1]
  [ 4  5]]

 [[ 6  7]
  [10 11]]]

As user @hpaulj correctly pointed out, this may be what you were looking for:

idx0, idx1 = [[0], [1]], [0, 2]
result = a[idx0, idx1]
print(result)

Instead of a slice, using a nested list.

Grismar0
Related