python - 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 中对其进行矢量化？

谢谢。

附言在这个例子中，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. (a-a[0])/thresh 进行整数除法(假设 a 具有整数 dtype)以将值分组 thresh 宽。

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

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

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