如何使用numpy.where？

Question

我想找到像x = np.array([[1, 1, 1], [2, 2, 2]])这样的数组的索引，其中元素等于y = np.array([1, 1, 1])。所以我这样做了：

In: np.where(x == y)
Out: (array([0, 0, 0]), array([0, 1, 2]))

这是正确的答案。但我希望只获得索引0，因为x的零元素等于y。

Answer 1

您需要先使用(x == y).all(axis=1)来比较axis=1的比较结果，即所有元素都相等：

np.where((x == y).all(axis=1))[0]
# array([0])