Search code examples
pythonnumpymatrix-indexing

selecting rows from 2 matrices using an indicator vector in numpy

Suppose I have a vector:

f = np.array([1,1,0,0]) #(4,)

and 2 matrices:

m1 = np.array([[1,2],[3,4],[5,6],[7,8]]) #(4,2)
m2 = np.array([[10,20],[30,40],[50,60],[70,80]]) #(4,2)

How can I create a new matrix m3 that selects rows from m1 where f == 1 and m2 otherwise?

I want m3 to be:

>>> m3
array([[ 1,  2],
       [ 3,  4],
       [50, 60],
       [70, 80]])

How do I achieve this? Would prefer a solution that I can use in theano as well.

Solution

  • I don't know about Theano but for numpy:

    np.where(f[:, None], m1, m2)