CODEWITHSUNDEEP
pythonarrayspython-3.xnumpy
lezaf
lezaf

Reputation: 512

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

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

Upvotes: 1

Views: 154

Answers (2)

Talha Tayyab
Talha Tayyab

Reputation: 27385

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.

Upvotes: 1

mozway
mozway

Reputation: 260640

You can use:

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

Output:

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

Upvotes: 1

Related Questions