numpy使用where函数定位满足条件的下标

官网DOC:https://numpy.org/doc/stable/reference/generated/numpy.where.html

np.where一般有两个用法:

  1. 当成if else的简写形式
  2. 得到满足条件的下标,方便后续操作

示例

当成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)

猜你喜欢

转载自blog.csdn.net/weixin_35757704/article/details/121175646