python库numpy使用技巧（一）——提取数组中非零元素

使用库numpy

Matlab中

• 通过逻辑矩阵可快速得到

a = [1,2,3,4,5,6]

a =

   1   2   3   4   5   6

b = logical([1,0,0,1,1,0])

b =

  1  0  0  1  1  0

a(b)

ans =

   1   4   5

为了在python中达到相同效果！！

Python中

• 导入numpy

import numpy as np

• 创建一个23的int型矩阵和一个23的bool型矩阵

a = np.array([[1,2,3],[4,5,6]])
b = np.array([[True,True,False],[False,True,False]])

• 矩阵直观显示

[[1 2 3]
 [4 5 6]]
[[ True  True False]
 [False  True False]]

• a, b两矩阵相乘

c = a*b

• 结果

[[1 2 0]
 [0 5 0]]

numpy.flatnonzero可返回a的展平版本中非零的索引。

np.flatnonzero(c)

• 结果

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

使用非零元素的索引作为索引数组来提取这些元素

c.ravel()[np.flatnonzero(c)]

• 结果

[1 2 5]

完整过程

import numpy as np

a = np.array([[1,2,3],[4,5,6]])
b = np.array([[True,True,False],[False,True,False]])
c = a*b
print(c.ravel()[np.flatnonzero(c)])

out:
[1 2 5]