在N维numpy数组中获取值的第一次出现

Question

我看到过这个问题，但我想要从mask = array == value生成的数组中进行简化。

mask = array([[[ True,  True,  True],
               [False,  True,  True]],

              [[False,  True,  True],
               [False,  True,  True]],

              [[False, False,  True],
               [False,  True,  True]]])

这会得到

where(mask) = (array([0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2]),
               array([0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1]),
               array([0, 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2]))

我想把它简化为一个只包含True第一次出现位置的数组。

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

但是我不知道怎么从numpy.where的输出中实现这一点。有人能帮帮我吗？

Answer 1

其实，这个问题很简单：

np.argmax(mask, 2)

举个例子：

In [15]: %paste
mask = array([[[ True,  True,  True],
               [False,  True,  True]],

              [[False,  True,  True],
               [False,  True,  True]],

              [[False, False,  True],
               [False,  True,  True]]])

## -- End pasted text --

In [16]: np.argmax(mask, 2)
Out[16]:
array([[0, 1],
       [1, 1],
       [2, 1]], dtype=int64)