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行的哪一列中出现了a的第j个元素。

结果将会是：

[1, 0, 1, 2]

Numpy是否提供这样的函数？

（感谢使用列表推导式回答，但从性能角度考虑，这不是一个选项。我也为在最后一个问题中提到Numpy而道歉。）

- Michael

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

- Benjamin

Answer 2

对于每一行的第一个匹配，可以像@Benjamin的帖子中所做的那样，在将a扩展为2D后使用argmax是一种高效的方法。

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

样例运行 -

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

- Divakar

Answer 3

不需要任何导入的懒惰解决方案：

a = [1, 5, 9, 7]

M = [
[0, 1, 9],
[5, 3, 8],
[3, 9, 0],
[0, 1, 7],
]

for n, i in enumerate(M):
    for j in a:
        if j in i:
            print("{} found at row {} column: {}".format(j, n, i.index(j)))

返回：

1 found at row 0 column: 1
9 found at row 0 column: 2
5 found at row 1 column: 0
9 found at row 2 column: 1
1 found at row 3 column: 1
7 found at row 3 column: 2

- Gábor Erdős

Answer 4

也许像这样？

>>> [list(M[i,:]).index(a[i]) for i in range(len(a))]
[1, 0, 1, 2]

- BPL

Answer 5

[sub.index(val) if val in sub else -1 for sub, val in zip(M, a)]
# [1, 0, 1, 2]

- Julien Spronck