在NumPy数组中的多个索引处获取元素

Question

I am trying to access specific rows and columns of a NumPy array as the documentation explains but I think I am missing something.

I have the following array:

arr = np.random.randint(10, size=(6, 4))
array([[1, 9, 6, 4],
       [8, 5, 0, 3],
       [3, 7, 3, 2],
       [1, 4, 8, 0],
       [5, 5, 8, 0],
       [0, 6, 4, 9]])

And I want to get the first and last row; and the first, third and last column, so I am trying:

arr[(0, -1),(0, 1, 3)]

But this is producing the following error:

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,) (3,)

I think I am misunderstanding this type of integer indexing. I would expect this output:

array([[1, 9, 4],
       [0, 6, 9]])

I can do it this way but it feels really awkward:

arr[(0,-1),:][:,(0,1,3)]

How can I get the i-th elements of different dimensions?

Answer 1

You are looking for np.ix_ :

arr[np.ix_((0, -1),(0, 1, 3))]

Results in

array([[1, 9, 4],
       [0, 6, 9]])

Answer 2

One way is to explicitly specify the index of each element of the result, where the key difference is that the row specification should be 2d:

>>> np.random.seed(444)
>>> arr = np.random.randint(10, size=(6, 4))
>>> arr
array([[3, 0, 7, 8],
       [3, 4, 7, 6],
       [8, 9, 2, 2],
       [2, 0, 3, 8],
       [0, 6, 6, 0],
       [3, 0, 6, 7]])

>>> rows = [[0],
...         [-1]]
>>> cols = [0, 1, -1]
>>> arr[rows, cols]
array([[3, 0, 8],
       [3, 0, 7]])

This example takes advantage of broadcasting for the row index. In other words, you can take advantage of the fact that the rows are repeating themselves for the remaining two columns in addition to the 0th.

To use advanced indexing one needs to select all elements explicitly. ... However, since the indexing arrays above just repeat themselves, broadcasting can be used. An example where you could not take advantage of broadcasting, and would need to specify each individual index, would be something like:

rows = [[1, 2],
        [3, 4]]
cols = [[2, 3],
        [1, 0]]

In that case, you're selecting elements at (1, 2), (2, 3), ...

See this example from the docs.

You might notice that this is exactly what is produced by np.ix_ (though I'm not saying it's not a useful function):

>>> rows, cols = np.ix_((0, -1),(0, 1, 3))