更改顺序 np.argmax 索引取自

Question

我有三个这样的矩阵：

one = np.empty((5,5))
one[:] = 10
two = np.empty((5,5))
two[:] = 10 
three = np.empty((5,5))
three[:] = 2

然后我把它们堆叠起来：

stacked = np.dstack([one, two, three])

最后确定具有最大值的索引：

t_max = np.argmax(stacked, axis=2)

我现在想确定最大值，但有一个警告。 如果有多个深度具有相同的最大值(如我的示例所示(，我想从最大深度返回索引。

正如现在t_max返回：

[[0 0 0 0 0]
[0 0 0 0 0]
[0 0 0 0 0]
[0 0 0 0 0]
[0 0 0 0 0]]

但我想回来：

[[1 1 1 1 1]
[1 1 1 1 1]
[1 1 1 1 1]
[1 1 1 1 1]
[1 1 1 1 1]]

由于第二个深度具有与第一个深度相同的最大值，但也具有更大的深度。

编辑：

我想我可以先做np.flip(stacked, axis=2)，但也许这不是最好的方法。

Answer 1

出于效率目的，我建议使用flipped视图，然后在减去最后一个轴长度后获取索引，如下所示 -

stacked.shape[-1] - stacked[...,::-1].argmax(-1) - 1

另一种不翻转和更长的方法，是与最大值进行比较，使用累积总和，然后使用argmax来捕获这些匹配项的最后一项，如下所示 -

(stacked == stacked.max(-1,keepdims=1)).cumsum(-1).argmax(-1)

示例运行 -

In [29]: stacked = np.random.randint(0,3,(2,5,3))
In [30]: stacked
Out[30]: 
array([[[2, 1, 2],
[1, 1, 1],
[1, 1, 1],
[1, 1, 0],
[1, 2, 2]],
[[2, 1, 1],
[1, 1, 1],
[1, 2, 2],
[1, 1, 0],
[1, 0, 0]]])
In [31]: stacked.shape[-1] - stacked[...,::-1].argmax(-1) - 1
Out[31]: 
array([[2, 2, 2, 1, 2],
[0, 2, 2, 1, 0]])
In [32]: (stacked == stacked.max(-1,keepdims=1)).cumsum(-1).argmax(-1)
Out[32]: 
array([[2, 2, 2, 1, 2],
[0, 2, 2, 1, 0]])

运行时测试 -

In [33]: stacked = np.random.randint(0,10,(1000,1000,100))
In [34]: %timeit stacked.shape[-1] - stacked[...,::-1].argmax(-1) - 1
1 loop, best of 3: 281 ms per loop
In [35]: %timeit (stacked == stacked.max(-1,keepdims=1)).cumsum(-1).argmax(-1)
1 loop, best of 3: 659 ms per loop