Getting elements at multiple indices in a NumPy array

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?

hilberts_drinking_problem · Accepted Answer · 2018-06-12 13:01:50Z

2

You are looking for np.ix_:

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

Results in

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

answered Jun 12, 2018 at 13:01

hilberts_drinking_problem

11.6k3 gold badges25 silver badges55 bronze badges

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

Comments

Brad Solomon · Accepted Answer · 2018-06-12 13:04:19Z

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))

Collectives™ on Stack Overflow

Getting elements at multiple indices in a NumPy array

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related