使用numpy.where()返回完整数组的索引，条件测试基于切片数组

Question

我有一个3 x 3 x 3的numpy数组，叫做a（后面的注释在你读完问题后就能理解了）：

array([[[8, 1, 0],     # irrelevant 1 (is at position 1 rather than 0)
        [1, 7, 5],     # the 1 on this line is what I am after!
        [1, 4, 9]],    # irrelevant 1 (out of the "cross")

       [[4, 0, 1],     # irrelevant 1 (is at position 2 rather than 0)
        [1, 0, 1],     # I'm only after the first 1 on this line!
        [6, 2, 1]],    # irrelevant 1 (is at position 2 rather than 0)

       [[0, 2, 2],
        [0, 6, 7],
        [3, 4, 9]]])

此外，我还有一个索引列表，指向这个矩阵的“中心十字”，叫做idx。

[array([0, 1, 1, 1, 2]), array([1, 0, 1, 2, 1])]

编辑：我之所以称它为“十字”，是因为它标记了下面的中心列和行：

>>> a[..., 0]
array([[8, 1, 1],
       [4, 1, 6],
       [0, 0, 3]])

我想要得到的是所有位于idx的数组中，第一项值为1的索引，但我在理解如何正确使用numpy.where()时遇到了困难。因为……

>>> a[..., 0][idx]
array([1, 4, 1, 6, 0])

……我尝试过……

>>> np.where(a[..., 0][idx] == 1)
(array([0, 2]),)

……但正如你所看到的，它返回的是切片数组的索引，而不是a的索引，而我想要的是：

[array([0, 1]), array([1, 1])]  #as a[0, 1, 0] and a [1, 1, 0] are equal to 1.

提前感谢你的帮助！

PS：在评论中，有人建议我尝试给出一个更广泛的适用场景。虽然这不是我正在使用的情况，但我想这可以用来处理图像，就像许多2D库那样，有一个源层、一个目标层和一个掩码（例如，参见cairo）。在这种情况下，掩码就是idx数组，人们可以想象处理RGB颜色的R通道（a[..., 0]）。

Answer 1

你可以使用 idx 来将索引转换回来：

>>> w = np.where(a[..., 0][idx] == 1)[0]
>>> array(idx).T[w]
array([[0, 1],
       [1, 1]])