官网DOC:https://numpy.org/doc/stable/reference/generated/numpy.where.html
np.where
一般有两个用法:
- 当成
if else
的简写形式 - 得到满足条件的下标,方便后续操作
示例
当成if else
的简写形式:
import numpy as np
my_array = np.array([1, 2, 3, 4, 5])
value = np.where(my_array > 3, True, False) # 条件为真就返回第二个参数True,为假就返回第三个参数False
print(value) # [False False False True True]
得到满足条件的下标:
my_matrix = np.array([
[1, 2, 3, 4],
[2, 2, 3, 4],
[3, 2, 3, 4],
])
value2 = np.where(my_matrix > 2)
print(value2)
# (array([0, 0, 1, 1, 2, 2, 2]), array([2, 3, 2, 3, 0, 2, 3]))
# 这里第1个array表示第一个维度的坐标
# 这里第2个array表示第二个维度的坐标
# 所以value2表示下标为[0,2],[0,3](这是my_matrix的第一行1234),[1,2],[1,3](这是my_matrix的第二行2234),[2,0],[2,2],[2,3](这是my_matrix的第三行3234)