根據一維數組獲取二維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 是二維數組的寬度。