如何根据Numpy数组a的元素值筛选保留数组b的指定元素？

给定两个NumPy数组a和b，要提取b中与a对应位置元素为1的元素，可通过以下步骤实现：

1. 先导入NumPy并定义目标数组：

import numpy as np

a = np.array([[0, 0, 0, 1, 1, 1, 0], [0, 0, 1, 1, 1, 0, 0]])
b = np.array([[5, 2, 0, 1, 0, 2, 3], [0, 0, 0, 1, 5, 3, 0]])

2. 使用布尔掩码+行遍历的方式筛选元素：

# 遍历每行，用a的对应行生成掩码筛选b的元素，转成列表匹配预期格式
result = [row[a_row == 1].tolist() for row, a_row in zip(b, a)]
print(result)

执行后输出：

[[1, 0, 2], [0, 1, 5]]

原理说明 #

• a_row == 1会生成与行长度一致的布尔数组，标记出a中值为1的位置；
• 用该布尔数组作为索引，就能从b的对应行中精准筛选出符合条件的元素；
• 通过列表推导式遍历所有行，最后将每行的NumPy数组转成普通列表，得到需求的格式。

如果需要保留NumPy数组格式，去掉.tolist()即可：

result_arr = np.array([row[a_row == 1] for row, a_row in zip(b, a)])
print(result_arr)
# 输出:
# [[1 0 2]
#  [0 1 5]]

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