如何获取nd.array中满足条件元素的一维nd.array类型索引而非tuple

问题根因 #

np.where 的返回值固定为元组结构：对n维输入数组，返回长度等于n的元组，每一个元素是对应维度的索引数组。

• 一维数组输入时，返回的元组仅包含1个元素，即实际需要的索引数组
• 直接对整个元组做np.array()转换，相当于把单元素数组按行堆叠，最终得到形状为(1, N)的二维数组，不符合1维返回值要求

正确实现方式 #

方法1：直接取np.where返回元组的第一个元素 #

这是最贴合原有代码逻辑的写法，仅需在原代码末尾加索引取值操作：

import numpy as np
arr = np.array([1, 4, 2, 1])
target_index = np.where(arr < 3)[0]

验证输出：

• 类型为np.ndarray
• 数组维度ndim=1
• 实际值为array([0, 2, 3])，完全匹配预期

方法2：使用np.flatnonzero #

这个方法天生设计为返回展平后的1维非零索引数组，不需要额外拆解元组，输出天然符合1维要求：

target_index = np.flatnonzero(arr < 3)

输出结果和方法1完全一致。

方法3：布尔掩码配合长度索引 #

通过生成全序列索引再用布尔值筛选，也能得到1维结果：

target_index = np.arange(arr.shape[0])[arr < 3]

补充提示：如果坚持通过转换整个np.where的返回结果来实现，需要对结果做展平操作（比如加.ravel()），但相比直接取元组元素的写法属于绕远路，没有必要。

内容的提问来源于stack exchange，提问作者james54