CODEWITHSUNDEEP
pythonarraysnumpyindexing
nullgeppetto
nullgeppetto

Reputation: 569

Matrix row-wise indexing

I have a numpy array, e.g., the following 

import numpy as np
A = np.array([[1, 2, 3], [4, 5, 6]])

and also another numpy array with boolean values, e.g., 

I = np.array([[True, False, False], [False, True, False]])

I would like to get the matrix whose elements' indices are given by I. In the above example, I'd like to get the matrix

array([[1], [5]])

but if I try

B = A[I]

then I get

array([1, 5])

I understand that this is due to the fact that the number of Trues's may not be the same in each row. But what about if they are? Is there any way of doing this using numpy?

In fact, I'd like to use this in Theano, using the tensor module. I have a theano expressions for the above (two T.matrix theano variables) that contain the above arrays. Is there any convenient way of computing the new, smaller matrix?

Upvotes: 1

Views: 95

Answers (1)

cs95
cs95

Reputation: 403050

If you can figure out how many items are returned from each row in advance, you can just reshape your output. I'd do it like this:

n = I.sum(1).max()
x = A[I].reshape(-1, n)

print(x)
array([[1],
       [5]])

Upvotes: 2

Related Questions