numpy.nonzero 的用法

发表于 更新于 分类于 阅读次数：

本文给出了 np.nonzero 的常见用法

返回非零元的索引

如：

1
2
3
4
5
6
7
>>> x = np.array([[1, 2, 3], [1, 0, 2], [3, 2, 1]])
>>> x
array([[1, 2, 3],
       [1, 0, 2],
       [3, 2, 1]])
>>> np.nonzero(x)
>>> (array([0, 0, 0, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 2, 0, 1, 2]))

其中输出的两个数组分别代表行数和列数，这里 (0,0), (0, 1), (0, 2) ... 不为零。

也可以用 transpose(nonzero(x)) 来按元素返回：

1
2
3
4
5
6
7
8
9
>>> np.transpose(np.nonzero(x))
array([[0, 0],
       [0, 1],
       [0, 2],
       [1, 0],
       [1, 2],
       [2, 0],
       [2, 1],
       [2, 2]])

于是，我们就可以用 x[nonzero(x)] 来提取数组中的非零元：

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

返回符合某条件的元素索引：

1
2
3
4
5
6
>>> x > 1
array([[False,  True,  True],
       [False, False,  True],
       [ True,  True, False]], dtype=bool)
>>> np.nonzero(x > 1)
(array([0, 0, 1, 2, 2]), array([1, 2, 2, 0, 1]))

MATLAB 中类似的操作：

1
2
3
4
5
6
7
8
9
10
11
>> x = [1 2 3; 1 0 2; 3 2 1]

x =

     1     2     3
     1     0     2
     3     2     1
>> x ~= 0
    1   1   1
    1   0   1
    1   1   1