numpy.where() 函数用法详解：条件选择与坐标提取

木白

发布时间：2026-10-13| 6人浏览过

原创

本文详细解析 numpy.where() 函数的两种核心用法。第一种形式 np.where(condition, x, y) 用于根据条件在 x 和 y 中进行元素级选择，类似列表推导式；第二种形式 np.where(condition) 仅返回满足条件元素的坐标元组，等价于 nonzero。文章涵盖一维、多维数组的处理及多条件组合查询技巧，强调括号使用的必要性。

numpy.where() 函数概述

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

在 NumPy 库中，numpy.where() 函数主要提供两种不同的调用模式：

三参数模式：条件选择

1. np.where(condition, x, y)

该模式用于根据布尔条件 condition 进行元素选择：当条件为真时输出对应位置的 x 值，否则输出 y 值。

对于一维数组的情况，其逻辑等同于 [xv if c else yv for (c,xv,yv) in zip(condition,x,y)]


>>> aa = np.arange(10)
>>> np.where(aa,1,-1)
array([-1, 1, 1, 1, 1, 1, 1, 1, 1, 1]) # 0为False，所以第一个输出-1
>>> np.where(aa > 5,1,-1)
array([-1, -1, -1, -1, -1, -1, 1, 1, 1, 1])

>>> np.where([[True,False], [True,True]],  # 官网上的例子
       [[1,2], [3,4]],
       [[9,8], [7,6]])
array([[1, 8],
    [3, 4]])

以上示例中的条件数组为 [[True,False], [True,False]]，这决定了最终输出结果的四个元素取值。具体而言，第一个元素从数组 [1,9] 中选取，因对应条件为 True，故取 1；第二个元素从 [2,8] 中选取，因条件为 False，故取 8，后续元素依此类推。可参考以下补充示例以加深理解：


>>> a = 10
>>> np.where([[a > 5,a < 5], [a == 10,a == 7]],
       [["chosen","not chosen"], ["chosen","not chosen"]],
       [["not chosen","chosen"], ["not chosen","chosen"]])

array([['chosen', 'chosen'],
    ['chosen', 'chosen']], dtype='<U10')

单参数模式：坐标提取

2. np.where(condition)

当仅传入条件参数 condition 而省略 x 和 y 时，函数将返回满足条件（即非零值）元素的坐标索引，其行为等价于 numpy.nonzero()。返回结果以元组（tuple）形式呈现，元组内包含的数组数量与原数组维度一致，分别对应各维度下符合条件元素的坐标。


>>> a = np.array([2,4,6,8,10])
>>> np.where(a > 5)       # 返回索引
(array([2, 3, 4]),)  
>>> a[np.where(a > 5)]       # 等价于 a[a>5]
array([ 6, 8, 10])

>>> np.where([[0, 1], [1, 0]])
(array([0, 1]), array([1, 0]))

在上述示例中，条件数组 [[0,1],[1,0]] 中有两个非零元素（值为1），它们的第一个维度坐标分别为 [0,1]，第二个维度坐标为 [1,0]。

接下来展示一个涉及更高维度的复杂示例：


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

    [[ 9, 10, 11],
    [12, 13, 14],
    [15, 16, 17]],

    [[18, 19, 20],
    [21, 22, 23],
    [24, 25, 26]]])

>>> np.where(a > 5)
(array([0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2]),
 array([2, 2, 2, 0, 0, 0, 1, 1, 1, 2, 2, 2, 0, 0, 0, 1, 1, 1, 2, 2, 2]),
 array([0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]))


# 符合条件的元素为
    [ 6, 7, 8]],

   [[ 9, 10, 11],
    [12, 13, 14],
    [15, 16, 17]],

   [[18, 19, 20],
    [21, 22, 23],
    [24, 25, 26]]]

由此可见，np.where() 会输出每个符合条件元素的完整坐标信息。由于原数组具有三维结构，因此返回的元组中包含三个数组，分别对应第三维、第二维和第一维的索引。