python - Numpy:在每行上查找元素的列索引

Question

假设我有一个包含要查找元素的向量:

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

现在我有一个矩阵，应该在其中搜索这些元素:

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

现在我想得到一个索引数组，告诉 M 的第 j 行的哪一列 的元素 j a 发生。

结果是:

[1, 0, 1, 2]

Numpy有提供这样的功能吗？

(感谢列表推导式的回答，但这不是性能方面的选择。我也很抱歉在最后一个问题中提到 Numpy。)

Answer 1

注意结果:

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

可以通过以下方式检索索引:

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