是否有一个 numpy 函数可以根据切片或范围中的值（而不是索引）对数组进行子集化？

Question

下面的代码可以工作，但在第二种情况下效率非常低，因为

i

被转换为完整的数组，增加了内存使用 和 运行时间。

import numpy as np

arr = np.array([0, 10, 20, 30, 40, 50, 29999999])
M = np.max(arr)

# slice based on indices
s = slice(2, None)
print(arr[s])

# slice based on values
s = slice(20, None)
i = range(*s.indices(M + 1))
print(arr[np.isin(arr, i)])  # inefficient!

输出：

[      20       30       40       50 29999999]
[      20       30       40       50 29999999]

是否有 numpy 函数可以直接改进它？我是否应该使用

np.where

代替使用切片的回调（感觉它也可能很慢，对于每个元素从 C++ 返回到 Python [或者不会这样做？]）？还是我错过了更好的方法？

Answer 1

如果我正确理解你的问题 - 你希望提取所有值，包括

20

和之后。在这种情况下，您可以利用累积总和，如下所示：

>>> arr[(arr==20).cumsum()>0] 
array([      20,       30,       40,       50, 29999999])

布尔值总和将返回 1/0，因此

cumsum

不能为负数，一旦为正数一次 - 则意味着遇到了

20

。