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