獲取 Python 矩陣數組中值的 (i,j) 索引的最快方法？

Question

在下面的 numpy 數組中獲取值 1 的所有 (i,j) 索引的最快方法是什么？

import numpy as np 

arr = [[1, 0, 1], [0, 0, 0], [0, 1, 0]]
arr = np.array(arr)
arr
array([[1, 0, 1],
       [0, 0, 0],
       [0, 1, 0]])

我曾嘗試for i for j使用嵌套 for 循環，但對於較大的數據集，這需要O(n^2) 。 我也嘗試過使用np.where(arr == 1)但沒有得到 (i,j) 索引：

np.where(arr == 1)
(array([0, 0, 2]), array([0, 2, 1])) # not (i,j) indices

獲取 (i,j) 索引的更好或更快的方法是什么？

Answer 1

您返回的兩個 arrays 只是每個匹配項的列/行索引。 您可以簡單地將 zip 第一個和第二個數組一起放入元組。

import numpy as np 

arr = [[1, 0, 1], [0, 0, 0], [0, 1, 0]]
arr = np.array(arr)    

d = np.where(arr == 1)
tuple(zip(d[0], d[1]))

Output

((0, 0), (0, 2), (2, 1))

Answer 2

如果您堅持將索引作為 ij 對，請使用np.argwhere ：

ij = np.argwhere(arr == 1)

ij將是一個 Nx2 數組，每行包含一個索引對。 您仍然可以將這些列用作精美的索引。

否則，請使用np.nonzero ，當您使用單個參數執行np.where時會調用它：

i, j = np.nonzero(arr)

如果您的數組嚴格為零或一，則在任何一種情況下都不需要== 0 。 這兩個函數都尋找非零元素。

另外，請記住，所有這些函數都具有相似的時間復雜度，即O(N) ，因為N是數組中元素的總數。 無論您使用 numpy function 還是嵌套for循環，都至少需要對數組進行一次遍歷。

Answer 3

所以這基本上是一個簡單的數組遍歷問題，這個問題總是需要 O(n^2) 時間。

所以我覺得這里可以使用循環

輸入：

arr = [[1, 0, 1], [0, 0, 0], [0, 1, 0]]
arr = np.array(arr)


def method_1(arr):
    temp = []
    for i in range(len(arr)):
        for j in range(len(arr[i])):
            if arr[i][j] == 1:
                temp.append((i, j))
    return temp


def method_2(arr):
    d= np.where(arr == 1)
    temp = []
    for i, j in zip(d[0], d[1]):
        temp.append((i, j))
    return temp


start_time = time.time()
method_1(arr)
end_time = time.time()

print(end_time- start_time)

start_time = time.time()
method_2(arr)
end_time = time.time()

print(end_time- start_time)

OUTPUT：

2.00271606445e-05          #Time taken by 1st algorithm
4.50611114502e-05          #Time taken by 2nd algorithm
[(0, 0), (0, 2), (2, 1)]
[(0, 0), (0, 2), (2, 1)]

注意：每當您想檢查算法的效率時，您可以使用時間 function