NumPy过滤一维数组后如何获取剩余元素对应原数组索引

NumPy 原生支持该场景需求，不需要手动匹配索引，有两种常用实现方式：

• 方案一：直接复用布尔掩码（推荐，性能最优）
  你编写的过滤条件 _nums['param1'] >= X 本身会返回一个和原 _nums 形状完全一致的布尔数组：值为True的位置代表元素满足保留条件，False代表元素会被过滤。不需要直接用该掩码索引_nums后丢弃掩码，同一个掩码可以直接同时索引两个同形状数组，得到的结果天然索引一一对应：

  # 第一步先单独生成过滤掩码
  filter_mask = _nums['param1'] >= X
  # 用同一个掩码分别过滤两个数组
  filtered_nums = _nums[filter_mask]
  filtered_data = _data[filter_mask]
  

  这是NumPy生态中处理同形状关联数组过滤的标准写法，没有额外的索引计算开销。
• 方案二：获取保留元素的原整数索引
  如果你确实需要拿到保留元素在原数组中的整数位置，可以用np.nonzero()（单参数调用时np.where()和它效果完全一致）处理布尔掩码，返回值就是满足条件元素的索引元组：

  filter_mask = _nums['param1'] >= X
  # 一维数组取返回元组的第一个元素，就是所有保留元素的原索引
  keep_index = np.nonzero(filter_mask)[0]
  # 用得到的索引分别取值即可
  filtered_nums = _nums[keep_index]
  filtered_data = _data[keep_index]
  

  以你给出的示例数据为例，最终得到的keep_index就是array([1, 4])，和你手动匹配的原数组位置完全一致。

注意：不要先对_nums做过滤再反查索引，过滤后的数组已经丢失了原位置信息，反查不仅性能差，遇到重复值还会出现匹配错误。

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