在numpy中找到两个数组中一对元素的位置

Question

我有两个 numpy 数组，分别为 x 和 y

假设 x = [0, 1, 1, 1, 3, 4, 5, 5, 5] 且 y = [0, 2, 3, 4, 2, 1, 3, 4, 5]

两个数组的长度相同，并且我确信要查找的坐标对在数组中一定存在。

如何在这些数组中找到 (a, b) 的索引位置，其中 a 是 x 中的一个元素，b 是对应的 y 中的元素。例如，(1, 4) 的索引位置将会是 3：在 x 和 y 的索引位置为 3 的元素分别为 1 和 4。

- shane

Answer 1

如果您想要一个纯粹的numpy解决方案，您可以使用numpy.where结合numpy.logical_and。

In [16]: import numpy as np

In [17]: x = np.array([0, 1, 1, 1, 3, 4, 5, 5, 5])

In [18]: y = np.array([0, 2, 3, 4, 2, 1, 3, 4, 5])

In [19]: np.where(np.logical_and(x == 1, y == 4))[0]
Out[19]: array([3], dtype=int64)

numpy.logical_and允许你在两个numpy数组之间逐元素执行逻辑AND操作。我们正在确定哪些位置同时包含x值为1和y值为4的相应位置。满足此条件的位置为True。 numpy.where确定该条件得到满足的数组位置。实际上，numpy.where分别作为两个元素的元组返回条件为True的行和列位置，但由于我们仅考虑一个维度，因此只有第一个元组是有效的，这就是为什么我们立即索引元组的第一个元素。

输出是一个numpy数组，其中包含条件成立的位置。您甚至可以进一步强制输出为索引列表，以使事情更加整洁/如果需要（感谢@EddoHintoso）：

In [20]: list(np.where(np.logical_and(x == 1, y == 4))[0])
Out[20]: [3]

- rayryeng

---

而 OP 可以通过包装 list() 将最后的输出强制转换为列表。这是目前为止最正确的答案。 - Eddo Hintoso

Answer 2

你可以将第一个数组与第一个值进行比较，将第二个数组与第二个值进行比较，然后找到两者都为True的位置。然后你可以使用argmax获取第一个True出现的索引：

x = np.array([0, 1, 1, 1, 3, 4, 5, 5, 5])
y = np.array([0, 2, 3, 4, 2, 1, 3, 4, 5])
idx = ((x == 1) & (y == 4)).argmax()

In [35]: idx
Out[35]: 3

In [36]: x == 1
Out[36]: array([False,  True,  True,  True, False, False, False, False, False], dtype=bool)

In [37]: y == 4
Out[37]: array([False, False, False,  True, False, False, False,  True, False], dtype=bool)

如果您可以有多个出现，则可以使用以下方式使用 非零：

idx_list = ((x == 1) & (y == 4))
idx = idx_list.nonzero()[0]

In [51]: idx
Out[51]: array([3], dtype=int64)

或者，如果您需要索引列表：

<code>In [57]: idx_list.nonzero()[0].tolist()
Out[57]: [3]
</code>

您可以用一行代码实现这个功能：

idx = ((x == 1) & (y == 4)).nonzero()[0]

- Anton Protopopov

---

如果(1,4)有多个可能的匹配，该怎么办？ - rayryeng

@rayryeng 那么你可以使用 nonzero，为此进行编辑。 - Anton Protopopov

好的。我投了一票。谢谢。 - rayryeng

Answer 3

x = [0, 1, 1, 1, 3, 4, 5, 5, 5]
y = [0, 2, 3, 4, 2, 1, 3, 4, 5]
w=zip(x,y)
w.index((1,4))

- Pavan

---

请使用NumPy解决。 - shane