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

Question

我有一个numpy数组，带有一堆单调递增的值。 说，

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

还说

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中进行矢量化吗？

谢谢。

PS在这个例子中，select_idx（a，thresh）= [0,4,5,8]

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

Answer 1

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

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

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

1. (aa[0])/thresh整数除法（假设a具有整数dtype）以将值分组为thresh wide。

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

3. 最后，我们放弃了最后一个索引，它对应的大小a 。 或者，如果索引的顺序无关紧要，您可以利用它来获取零索引：

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