How to slice a numpy array using index arrays with different shapes?

Question

Let’s say that we have the following 2d numpy array:

arr = np.array([[1,1,0,1,1],
                [0,0,0,1,0],
                [1,0,0,0,0],
                [0,0,1,0,0],
                [0,1,0,0,0]])

and the following indices for rows and columns:

rows = np.array([0,2,4])
cols = np.array([1,2])

The objective is to slice arr using rows and cols to take the following expected result:

arr_sliced = np.array([[1,0],
                       [0,0],
                       [1,0]])

Using directly the arrays as indices like arr[rows, cols] leads to:

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

So what is the straightforward way to achieve this kind of slicing?

Update: useful information about the solution

So the solution was simple enough and it demands a basic comprehension about numpy’s broadcasting. Someone could read these nice but not so representative examples from numpy. Also, the general broadcasting rules explains why there is no shape mismatch in:

arr[rows[:, np.newaxis], cols]
# rows[:, np.newaxis].shape == (3,1)
# cols.shape == (2,)

Asked By: lezaf

||

Source

Answer 1

It looks like it is much quicker than indexing for large arrays.

arr[np.ix_([0,2,4],[1,2])]

array([[1, 0],
       [0, 0],
       [1, 0]])

document: https://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.ix_.html

This function takes N 1-D sequences and returns N outputs with N dimensions each, such that the shape is 1 in all but one dimension and the dimension with the non-unit shape value cycles through all N dimensions.

Answered By: God Is One

Answer 2

You can use:

arr[rows[:,None], cols[None]]

Output:

array([[1, 0],
       [0, 0],
       [1, 0]])

Answered By: mozway

How to slice a numpy array using index arrays with different shapes?

Question:

Update: useful information about the solution

Answers: