Numpy逐元素操作

Question

假设我有一个长度为n的列向量y，还有一个大小为n*m的矩阵X。我想检查y中的每个元素i是否在X中对应的行中出现。最有效的方法是什么？

例如： y = [1,2,3,4].T

和

X =[[1, 2, 3],[3, 4, 5],[4, 3, 2],[2, 2, 2]]

那么输出应该是：

[1, 0, 1, 0] or [True, False, True, False] 

无论哪种都行。

当然，我们可以使用for循环来迭代y和X，但还有更高效的方法吗？

- zycuber

---

你确定不想要输出为 np.array([True, False, True, False]) 吗？ - Eric

Answer 1

使用broadcasting的向量化方法 -

((X == y[:,None]).any(1)).astype(int)

样例运行 -

In [41]: X        # Input 1
Out[41]: 
array([[1, 2, 3],
       [3, 4, 5],
       [4, 3, 2],
       [2, 2, 2]])

In [42]: y        # Input 2
Out[42]: array([1, 2, 3, 4])

In [43]: X == y[:,None] # Broadcasted  comparison
Out[43]: 
array([[ True, False, False],
       [False, False, False],
       [False,  True, False],
       [False, False, False]], dtype=bool)

In [44]: (X == y[:,None]).any(1) # Check for any match along each row
Out[44]: array([ True, False,  True, False], dtype=bool)

In [45]: ((X == y[:,None]).any(1)).astype(int) # Convert to 1s and 0s
Out[45]: array([1, 0, 1, 0])

- Divakar

---

我认为any(1)应该写成.any(axis=1)以增加清晰度。 - Eric