Numpy: Find column index for element on each row

Question

Suppose I have a vector with elements to find:

a = np.array([1, 5, 9, 7])

Now I have a matrix where those elements should be searched:

M = np.array([
[0, 1, 9],
[5, 3, 8],
[3, 9, 0],
[0, 1, 7]
])

Now I'd like to get an index array telling in which column of row j of M the element j of a occurs.

The result would be:

[1, 0, 1, 2]

Does Numpy offer such a function?

(Thanks for the answers with list comprehensions, but that's not an option performance-wise. I also apologize for mentioning Numpy just in the final question.)

Answer 1

Note the result of:

M == a[:, None]
>>> array([[False,  True, False],
           [ True, False, False],
           [False,  True, False],
           [False, False,  True]], dtype=bool)

The indices can be retrieved with:

yind, xind = numpy.where(M == a[:, None])
>>> (array([0, 1, 2, 3], dtype=int64), array([1, 0, 1, 2], dtype=int64))

Answer 2

For the first match in each row, it might be an efficient way to use argmax after extending a to 2D as done in @Benjamin's post -

(M == a[:,None]).argmax(1)

Sample run -

In [16]: M
Out[16]: 
array([[0, 1, 9],
       [5, 3, 8],
       [3, 9, 0],
       [0, 1, 7]])

In [17]: a
Out[17]: array([1, 5, 9, 7])

In [18]: a[:,None]
Out[18]: 
array([[1],
       [5],
       [9],
       [7]])

In [19]: (M == a[:,None]).argmax(1)
Out[19]: array([1, 0, 1, 2])