Python中的 numpy.flatnonzero()

numpy.flatnonzero()是一个用于返回数组中非零元素的索引数组的函数。

语法

numpy.flatnonzero(arr)

参数：

• arr：输入数组。

返回值：

• 返回值是一个一维数组，其中每个元素都是输入数组中非零元素的索引。

示例

import numpy as np

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

print(np.flatnonzero(arr))

输出：

[1 3 4 6 7]

解释：

• 输入数组共有5个非零元素，它们的索引分别为1，3，4，6，7。

应用

numpy.flatnonzero()函数可以用于获取数组中非零元素的位置，便于我们进行相应的处理。例如，将非零元素赋值为1，其他元素赋值为0，可以使用以下代码实现：

import numpy as np

arr = np.array([1, 0, 2, 0, 3, 0])
idx = np.flatnonzero(arr)

new_arr = np.zeros(arr.shape)
new_arr[idx] = 1

print(new_arr)

输出：

[1. 0. 1. 0. 1. 0.]

解释：将非零元素的位置赋值为1，即可得到一个新的数组。

注意事项

如果输入数组中无非零元素，则返回一个空数组。

import numpy as np

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

print(np.flatnonzero(arr))

输出：

[]