引用几个例子
>>> a = np.array([-1, -2, 0, 1, 2])
>>> np.where(a)
(array([0, 1, 3, 4]),)
只有input,后面没有指定条件,返回非0的元素
>>>a = np.arange(10)
>>>a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>>np.where(a < 5, a, 10*a)
array([ 0, 1, 2, 3, 4, 50, 60, 70, 80, 90])
有指定条件,等价为 if a的元素 < 5, 元素 =它本身, else 元素 = 10*它本身
>>> np.where([[True, False], [True, True]],
[[1, 2], [3, 4]],
[[9, 8], [7, 6]])
output:
array([[1, 8],
[3, 4]])
这里相当于一个2x2数组(T, F ; T, T), 其中" ; "表示换行,它后面跟的有两个2x2数组,分别对应了if T 条件 和else的条件。
具体来说,
如果第一行的(T, F)的元素为(T, T),它就等于(1, 2), 否则等于(9, 8)
现在第一个元素为T,就等于1, 第2个元素为F,就等于8。所以第一行返回(1, 8)
同理第二行。
再举个例子,
>>> b = np.array([[1,1],[2,3] ,[0.2,0.5]])
>>> b
array([[1. , 1. ],
[2. , 3. ],
[0.2, 0.5]])
>>> np.where(b)
output:
(array([0, 0, 1, 1, 2, 2]), array([0, 1, 0, 1, 0, 1]))
你肯定想说,为什么返回两个一维数组,这是什么意思。
首先,b是一个3x2数组,np.where要返回b中> 0的元素位置
位置包含了行坐标和列坐标,
那第一行就是行坐标,第2行就是列坐标,不信把它们拼在一起看一下
0, 0, 1, 1, 2, 2
0, 1, 0, 1, 0, 1
按列看,是不是坐标就是(0,0), (0, 1), (1, 0), … , (2,1)
把b改一下,把最后一个元素改为0
>>> b
array([[1. , 1. ],
[2. , 3. ],
[0.2, 0. ]])
>>> np.where(b)
output:
(array([0, 0, 1, 1, 2]), array([0, 1, 0, 1, 0]))
可以看到,最后一个元素的坐标没有了。