numpy array integer indexing in more than one dimension

Question

I'm pretty sure I'm missing something with integer indexing and could use some help. Say that I create a 2D array:

>>> import numpy as np
>>> x=np.array(range(24)).reshape((4,6))
>>> 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]])

I can then select row 1 and 2 with:

>>> x[[1,2],:]
array([[ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17]])

Or the column 1 of rows 2 and 3 with:

>>> x[[1,2],1]
array([ 7, 13])

So it would makes sense to me that I can select columns 3, 4 and 5 of rows 1 and 2 with this:

>>> x[[1,2],[3,4,5]]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: shape mismatch: objects cannot be broadcast to a single shape

And instead I need to do it in two steps:

>>> a=x[[1,2],:]
>>> a
array([[ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17]])
>>> a[:,[3,4,5]]
array([[ 9, 10, 11],
       [15, 16, 17]])

Coming from R, my expectations seem to be wrong. Can you confirm that this is indeed not possible in one step, or suggest a better alternative? Thanks!

EDIT: please note my choice of rows and columns in the example happen to be consecutive, but they don't have to be. In other words, slice indexing won't do for my case.

Answer 1

You also have the option of using broadcasting among the indexing arrays, which is what I would normally do, rather than indexing twice, which creates an intermediate copy of your data:

>>> x[[[1], [2]],[[3, 4, 5]]]
array([[ 9, 10, 11],
       [15, 16, 17]])

To see a little better what is going on and how to handle larger numbers of indices:

>>> row_idx = np.array([1, 2])
>>> col_idx = np.array([3, 4, 5])
>>> x[row_idx.reshape(-1, 1), col_idx]
array([[ 9, 10, 11],
       [15, 16, 17]])