根据一维数组获取二维numpy数组的索引

Question

我正在使用 numpy arrays 我遇到了障碍，也许你可以帮助我。 所以我有一个二维 numpy 数组，这个二维数组中的每个数组都有一个最大值，我需要的是获取最大值的索引，而不是二维数组的索引，而是像使用一维数组一样使用索引，我知道我可以使用numpy.argmax来查找索引，但是它的索引是根据二维数组而不是一维数组，也许举个例子会更好：

import numpy as np

arr =  np.array([[512, 523, 491], 
                 [512, 531, 495]])

index = np.argmax(arr, axis = 1)
index2 = np.argmax(arr)

print(index)
print(index2) 

index = [1, 1] 
index2 = 4

我的意思是它做了它应该做的事情，但是如果我需要像一维数组一样的索引怎么办？ 这是我想要的 output：

index = [1, 4] 

所以第一个最大值是 523，所以第一个索引是 1，第二个最大值是 531，所以第二个索引是 4，就像一维数组一样。 也许这是一个新手问题，但我不知道如何做到这一点，这只是一个例子， arr数组可以更大，所以，任何帮助将不胜感激，谢谢！

Answer 1

我能想到的最简单的事情就是在 index.html 中添加一个偏移量。

offset = np.arange(0,arr.shape[0]*arr.shape[1], arr.shape[1])
index1D = offset + index

对于上下文，np.arange 需要 3 个参数：开始、停止和跨步。 所以偏移量只是将特定行之前的元素总数添加到索引中。

Answer 2

第一行实际上是您想要的。 它返回每个嵌套数组中最大值的索引，即531在第二个嵌套数组的索引 1 中。

您可以将其与高级索引一起使用

>>> print(arr[np.arange(arr.shape[0]), index])
[523, 531]

或将索引更改为平面索引

flat_index = index + np.arange(0, arr.size, arr.shape[1])
# flat_index == [1, 4]

Answer 3

假设您知道二维数组的尺寸，您可以为每个元素添加偏移量。

例如

index = [1, 1]
offset = [0, 3]

index + offset
>>> [1, 4]

您可以通过执行np.arange(len(index))*size轻松获得偏移量。 其中 size 是二维数组的宽度。