numpy where函数

函数原型

numpy.where(condition[,x, y])

函数解释

根据条件返回x或者y中的元素，满足条件输出x中的值，不满足输出y中的值；如果只有条件，则输出满足条件对应的索引。

函数用法

>>> a = np.array([[1, 2, 3],
	      [4, 5, 6],
	      [7, 8, 9]])
>>> a
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
>>> b = np.array([[3, 1, 4],
	      [3, 5, 6],
	      [8, 2, 1]])
>>> b
array([[3, 1, 4],
       [3, 5, 6],
       [8, 2, 1]])
>>> c = np.array([[2, 4, 7],
	      [3, 2, 2],
	      [1, 4, 2]])
>>> c
array([[2, 4, 7],
       [3, 2, 2],
       [1, 4, 2]])
>>> d = np.where(a<5, b, c)
>>> d
array([[3, 1, 4],
       [3, 2, 2],
       [1, 4, 2]])
>>> e = np.where(a>2)
# 返回两个array，第一个array是从行上对符合条件的值进行描述的，
# 第二个array是从列上对符合条件的值进行描述的
# 比如数组上a的第一个符合条件的元素3，元素3在行上的索引为0，在列上的索引为2，
# 分别对应于两个array中的第一个位置的值，共有7个符合条件的值，所以两个array的大小为7
>>> e
(array([0, 1, 1, 1, 2, 2, 2], dtype=int64), array([2, 0, 1, 2, 0, 1, 2], dtype=int64))