[干货] 一文看懂numpy.nonzero() 与 numpy.argwhere()非零元素处理

0. 前言

在各类深度学习的过程中, 难免对非零元素进行处理.在Numpy中,提供了多种非零元素处理的接口和syntactic sugar. 其中就包括 numpy.nonzero() 与 numpy.argwhere()这两个函数,下面,本文将介绍这两个的使用方法与选择思想.

1. numpy.nonzero()

返回非零元素的索引。

返回一个数组元组，a 的每个维度一个，包含该维度中非零元素的索引。

例子:

test = np.array([0, 1, 0, 0, 0, 1, 0, 1, 0]).reshape(3, 3)
print(test)
print(np.nonzero(test))
print(len(np.nonzero(test)))

输出:

[[0 1 0]
 [0 0 1]
 [0 1 0]]
(array([0, 1, 2]), array([1, 2, 1]))
2

上一段代码首先生成了一个3x3的单位矩阵,非0元素的位置为[0,1],[1,2],[2,1]

同时,还有更多例子可以验证:

a.

x = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]])

x
array([[3, 0, 0],
       [0, 4, 0],
       [5, 6, 0]])

np.nonzero(x)
(array([0, 1, 2, 2]), array([0, 1, 0, 1]))

b.

x[np.nonzero(x)]
array([3, 4, 5, 6])

np.transpose(np.nonzero(x))
array([[0, 0],
       [1, 1],
       [2, 0],
       [2, 1]])

tips:

1. 如果对空数组操作,会返回一个len为1的tuple,元素为空array

2. 当在零维数组或标量上调用时，nonzero(a) = nonzero(atleast_1d(a))
   (1.17.0 版后已弃用：如果此行为是有意而为之的，请明确使用 atleast_1d。)

3. 要按元素而不是维度对索引进行分组，请使用 numpy.argwhere()，它为每个非零元素返回一行。

4. 虽然可以使用 a[nonzero(a)] 获得非零值，但建议使用 x[x.astype(bool)] 或 x[x != 0] 代替，这将正确处理 0-d 数组。
   转载请注明出处:https://blog.csdn.net/ftimes/article/details/119898004

2. numpy.argwhere()

输入一个数组,返回一个 (N, a.ndim)大小的索引数组.

以刚才那个例子为例:

test = np.array([0, 1, 0, 0, 0, 1, 0, 1, 0]).reshape(3, 3)
# test = np.array([])
print(test)
print(np.argwhere(test != 0))
print(np.argwhere(test != 0).shape)

输出:

[[0 1 0]
 [0 0 1]
 [0 1 0]]
[[0 1]
 [1 2]
 [2 1]]
(3, 2)

非0元素的位置为[0,1],[1,2],[2,1]
区别于numpy.nonzero() ,只返回一个 (N, a.ndim)大小的索引数组

同时, numpy.argwhere()还具有其他功能:

x = np.arange(6).reshape(2,3)

x
array([[0, 1, 2],
       [3, 4, 5]])

np.argwhere(x>1)
array([[0, 2],
       [1, 0],
       [1, 1],
       [1, 2]])

tips:

1. 在不写表达式时,即np.argwhere(test), 与 np.argwhere(test != 0)作用相一致,但推荐用此方法处理0d数组
2. 同时, np.argwhere拥有更多功能,例如:np.argwhere(x>1)
3. argwhere 的输出可能不适合索引数组。为此，可以改用 nonzero(a)。