python返回数组_python – 使用numpy.where()返回完整数组的索...

我有以下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]).