根据另一个数组的样本从 numpy 数组中获取样本

Question

我有一个矩阵a 和另一个向量b：

a = np.array([[0,1,2],
              [3,4,5],
              [6,7,8]]

b = np.array([0.5,0.3,0.2])

是否有任何速记函数可以根据b 的值从a 中提取行？

例如，从a 获取样本，b >= 0.3 应该返回：

[[0,1,2],
 [3,4,5]]

---

我想出了什么：

a_idx = np.argwhere(b >= 0.3).flatten()
c = np.take(a, a_idx)

这可行，但我想知道是否有更好的解决方案（特别适用于大型数组）。

提前致谢。

Answer 1

a[b >= 0.3] 更简洁，但您的方法似乎更快（在 1M 条目阵列上进行了 200 次测量测试，您的方法快了约 33%）。