numpy - 从带有间距的数组中选择元素

我有一个 numpy 数组，其中包含一堆单调递增的值。说，

a = [1,2,3,4,6,10,10,11,14]
a_arr=np.array(a)

Also say

thresh = 4

我想创建一个包含子集索引的数组a_arr它逐步遍历数组，选择元素但忽略至少没有间隔的元素thresh远离最后的选择。用算法可以更容易地描述这一点：

def select_idx(a, thresh):
    ret = []
    for idx, elt in enumerate(a):
        if len(ret) == 0 or elt >= a[ret[-1]] + thresh:
            ret.append(idx)
    return ret

显然我可以使用这个函数来做到这一点，但这看起来很慢。有什么方法可以在 numpy 中对其进行向量化吗？

Thanks.

附：在此示例中，select_idx(a, thresh) = [0, 4, 5, 8]

Edit：这个问题的近似版本可能更容易矢量化：将数轴划分为大小的桶thresh，我猜从 a 中的第一个值开始。因此，此示例中的存储桶分隔符将为 0、4、8、12、16...。选择作为存储桶中第一个元素的数字的索引。 （是的，我意识到这与我之前写的不一样。）

这是您的近似问题的矢量化解决方案：

idx = np.cumsum(np.bincount((a-a[0])/thresh))[:-1]

这将为您提供除始终存在的第一个零之外的所有索引。解释如下：

1. (a-a[0])/thresh进行整数除法（假设a有一个整数数据类型）将值分组thresh wide.

2. cumsum(bincount(...))计算每个组的大小并将其转换为索引。请注意，如果存储桶中没有值bincount将报告 0，因此该数组中可能存在重复。

3. 最后，我们丢弃最后一个索引，它对应于a。或者，如果索引的顺序无关紧要，您可以利用它来恢复零索引：

   idx = np.cumsum(np.bincount((a-a[0])/thresh)) % len(a)