使用 numpy 獲取每列特定值的行索引

Question

我有一個 10,000 x 10,000 的矩陣，其中填充了 1 和 0。 我想要做的是遍歷每一列並找到包含值 1 的行。

然后我想將它存儲在一個有 2 列的新矩陣中：第 1 列 = 列索引，第 2 列 = 包含 1 的行索引數組。有些列根本沒有任何 1，在這種情況下它會是一個空數組。

嘗試再次執行 for 循環，但計算效率低下。

我嘗試使用較小的矩陣

#sample matrix
n = 4
mat = [[randint(0,1) for _ in range(n)] for _ in range(n)]

arr = np.random.randint(0, size=(4, 2))

for col in range(n):
    arr[n][1] = n
    arr[n][2] = np.where(col == 1)

但是對於 10,000 x 10,000 矩陣，這運行得非常緩慢。 我想知道這是否正確，是否有更好的方法？

Answer 1

獲取索引，其中a[i][j] == 1

您可以使用numpy.argwhere()或numpy.nonzero()有效地獲取您正在尋找的數據（零和一矩陣中的位置numpy.nonzero() ，但是您將無法以指定的格式獲取它們在您單獨使用 NumPy ndarrays 的原始問題中。

您可以使用 ndarrays 和標准 Python 列表的組合以您指定的格式獲取數據，但是由於考慮到您正在使用的數據的大小，效率至關重要，我認為最好專注於獲取數據而不是獲取它采用不規則 Python 列表的 ndarray 格式。

如果您提到的格式是一個硬性要求，您總是可以在計算之后重新格式化結果（矩陣中的索引1 ），這樣您的代碼將受益於 NumPy 在繁重計算期間提供的優化 - 減少您的執行時間程序整體。

使用np.argwhere()示例

import numpy as np

a = np.random.randint(0, 2, size=(4,4))
b = np.argwhere(a == 1)

print(f'a\n{a}')
print(f'b\n{b}')

輸出

a
[[1 1 1 1]
 [0 0 0 0]
 [1 0 1 0]
 [1 1 1 1]]
b
[[0 0]
 [0 1]
 [0 2]
 [0 3]
 [2 0]
 [2 2]
 [3 0]
 [3 1]
 [3 2]
 [3 3]]

如您所見， np.argwhere(a == 1)返回一個 ndarray，其值為 ndarrays，其中包含a的位置索引，其值 ( x ) 滿足條件x == 1 。

我給了上面的方法a = np.random.randint(0, 2, size=(10000,10000)在我的筆記本電腦上嘗試了幾次（沒什么特別的），每次都在 3-5 秒左右完成。

獲取所有值都為!= 1行索引

如果你想存儲的所有行索引a不含值== 1 ，最簡單的方法（假設你使用我的示例代碼段）很可能是通過使用numpy.setdiff1d()返回行索引是數組內不存在b -即含有的所有的行索引的陣列之間的差集a和1D陣列b[0]這將是所有值的行索引在a是!= 1

假設a和b與上例相同。

c = np.setdiff1d(np.arange(a.shape[0]), b[:, 0])
print(c)

輸出

array([1])

在上面的例子c = [1]為1處於僅行索引a不包含任何值== 1 。

值得注意的是，如果a定義為np.random.randint(0, 2, size=(10000,10000) ，則c不是零長度（即空）數組的概率非常小。這是因為對於不包含值== 1 ， np.random必須連續返回0 10,000 次才能用0填充一行。

為什么要使用多個 NumPy 數組？

我知道使用b和c分別存儲與a == 1和a != 1位置有關的結果似乎很奇怪。 為什么不使用原始問題中概述的不規則list ？

簡而言之，答案是效率。 通過使用 NumPy 數組，您將能夠對數據進行矢量化計算，並在很大程度上避免代價高昂的 Python 循環，其好處將被大大放大，這反映在給定您正在處理的數據大小的執行時間上。

您始終可以以更人性化的不同格式存儲數據，並根據需要將其映射回 NumPy，但是與原始問題中的示例相比，上述示例可能會在執行時顯着提高效率。