Numpy 变更单np.argmax索引取自

我有三个矩阵：

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（堆叠，轴=2）

，但这可能不是最好的方法。

为了提高效率，我建议使用

翻转的

视图，然后从上一个轴长度中减去后得到索引，就像这样-

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

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

另一种不翻转且稍长的方法是，与最大值进行比较，使用累积求和，然后使用

argmax

捕获这些匹配中的最后一个，如下所示-

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

(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