python - 将 np.where 应用于方括号过滤以进行 numpy 过滤

Question

我可以通过 numpy 数组过滤

a[np.where(a[:,0]==some_expression)]

或

a[a[:,0]==some_expression]

每个版本的(缺点)优势是什么 - 特别是在性能方面？

Answer 1

bool 索引在内部转换为整数索引。这是 indicated in the docs :

In general if an index includes a Boolean array, the result will be identical to inserting obj.nonzero() into the same position and using the integer array indexing mechanism described above.

所以这两种方法的复杂度是一样的。但是 np.where 对于大型数组更有效:

np.random.seed(0)
a = np.random.randint(0, 10, (10**7, 1))
%timeit a[np.where(a[:, 0] == 5)]  # 50.1 ms per loop
%timeit a[a[:, 0] == 5]            # 62.6 ms per loop

现在 np.where 还有其他好处:高级整数索引在多个维度上都能很好地工作。有关 bool 索引在这方面不直观的示例，请参见 NumPy indexing: broadcasting with Boolean arrays .由于 np.where 比 bool 索引更有效，这只是它应该被首选的一个额外原因。