如何快速尋找NumPy數組中的非零元素位置

NumPy是Python中一個強大的科學計算庫，提供了高性能的多維數組和矩陣計算功能。在數據分析和處理中，經常需要對數組中的非零元素進行尋找和相關的計算操作。本文將討論使用NumPy庫中的函數來找到數組中的非零元素位置。

一、NumPy庫中尋找非零元素位置函數

NumPy提供了三個函數來尋找非零元素的位置：

numpy.nonzero(a)
numpy.where(condition, [x, y])
numpy.argwhere(a)

這些函數分別返回一個包含非零元素位置的元組。下面將詳細介紹每個函數的用途和示例。

二、numpy.nonzero(a)

numpy.nonzero(a)函數返回一個數組a中非零元素的索引。返回的結果是一個長度為a.ndim的元組，其中每個元素是一個數組，對應於一個維度。因此，返回的結果可以直接用於構建一個索引元組。

以下是使用numpy.nonzero(a)函數查找一維和二維數組中非零元素位置的示例：

import numpy as np

# 一維數組示例
a = np.array([1,0,2,0,3])
print(np.nonzero(a))

# 二維數組示例
b = np.array([[0, 1, 0], [1, 0, 1], [1, 1, 0]])
print(np.nonzero(b))

運行結果如下：

(array([0, 2, 4], dtype=int64),)
(array([0, 1, 1, 2, 2], dtype=int64), array([1, 0, 2, 0, 1], dtype=int64))

解釋一下以上代碼中的numpy.nonzero()函數返回的結果。

對於一維數組[1,0,2,0,3]，numpy.nonzero()函數返回的是(array([0, 2, 4], dtype=int64),)。這表示數組a的0、2、4位置上分別有非零元素。

對於二維數組[[0, 1, 0], [1, 0, 1], [1, 1, 0]]，numpy.nonzero()函數返回的是(array([0, 1, 1, 2, 2], dtype=int64), array([1, 0, 2, 0, 1], dtype=int64))。這表示數組b的(0,1)、(1,0)、(1,2)、(2,0)、(2,1)這5個位置上分別有非零元素。

三、numpy.where(condition, [x, y])

numpy.where(condition, [x, y])函數返回滿足指定條件的輸入數組元素的索引。此函數相當於if-else語句的向量化版本。如果只有滿足條件的元素的索引需要被提取，那麼使用此函數將更加高效。

以下是使用numpy.where(condition, [x, y])函數查找一維和二維數組中非零元素位置的示例：

import numpy as np

# 一維數組示例
a = np.array([1,0,2,0,3])
print(np.where(a!=0)[0])

# 二維數組示例
b = np.array([[0, 1, 0], [1, 0, 1], [1, 1, 0]])
print(np.where(b!=0))

運行結果如下：

[0 2 4]
(array([0, 1, 1, 2, 2], dtype=int64), array([1, 0, 2, 0, 1], dtype=int64))

解釋一下以上代碼中numpy.where(condition, [x, y])函數返回的結果。

對於一維數組[1,0,2,0,3]，np.where(a!=0)[0]將返回0, 2, 4的列表，這些是非零元素的索引。

對於二維數組[[0, 1, 0], [1, 0, 1], [1, 1, 0]]，np.where(b!=0)返回的是兩個長度相等的數組，第一個數組是非零元素所在的行，第二個數組是非零元素所在的列。

四、numpy.argwhere(a)

numpy.argwhere(a)函數返回一個數組中非零元素的坐標，即各維度上非零元素的索引列表。此函數等價於調用np.transpose(np.nonzero(a))，當然還有其他方式可以實現相同的功能。

以下是使用numpy.argwhere(a)函數查找一維和二維數組中非零元素位置的示例：

import numpy as np

# 一維數組示例
a = np.array([1,0,2,0,3])
print(np.argwhere(a))

# 二維數組示例
b = np.array([[0, 1, 0], [1, 0, 1], [1, 1, 0]])
print(np.argwhere(b))

運行結果如下：

[[0]
 [2]
 [4]]
[[0 1]
 [1 0]
 [1 2]
 [2 0]
 [2 1]]

解釋一下以上代碼中numpy.argwhere(a)函數返回的結果。

對於一維數組[1,0,2,0,3]，numpy.argwhere(a)返回的結果是一個列表，它展示了在a中非零元素所在的位置。

對於二維數組[[0, 1, 0], [1, 0, 1], [1, 1, 0]]，numpy.argwhere(b)返回的結果是一個包含數組的列表，每個數組中有兩個元素，它們對應於二維矩陣b中非零元素的行和列。

五、總結

各種不同的情況和數據類型下，尋找非零元素位置的函數有所不同，你需要比較它們的優缺點來決定使用哪個函數。儘管有多種函數可以選擇，但在應用中使用適當的函數可以使代碼更加清晰和高效。

希望通過本文的講解，你對如何快速尋找NumPy數組中的非零元素位置有了更加深刻的了解。