numpy where

 

>>> import numpy as np
>>> x = np.arange(9.).reshape(3, 3)
>>> x
array([[ 0.,  1.,  2.],
       [ 3.,  4.,  5.],
       [ 6.,  7.,  8.]])
>>> np.where( x > 5 )
(array([2, 2, 2]), array([0, 1, 2]))       第一个array返回的是第一维的坐标，第二个aray返回的是第二维的坐标。一二维的坐标结合，你就会发现就是6，7，8这三个数 
>>> x[np.where( x > 5 )]
array([ 6.,  7.,  8.])

 可以看到，where返回的都是坐标，要注意where对于一维array的返回

>>> a = np.array([1,2,3,4,5,6])
>>> b = np.where(a>3)
>>> b
(array([3, 4, 5]),)
>>> b[0]
array([3, 4, 5])
>>> a = np.array([1,3,5,6,7,8])
>>> b = np.where(a>2)
>>> b
(array([1, 2, 3, 4, 5]),)