如何获得至少两个连续值都大于阈值的索引？

Question

例如，让我们考虑以下 numpy 数组：

[1, 5, 0, 5, 4, 6, 1, -1, 5, 10]

另外，我们假设阈值等于3。也就是说，我们正在寻找至少两个连续值都高于阈值的序列。

输出将是这些值的索引，在我们的例子中是：

[[3, 4, 5], [8, 9]]

如果输出数组被展平，那也可以！

[3, 4, 5, 8, 9]

---

输出说明

在我们的初始数组中，我们可以看到 forindex = 1我们有 value 5，它大于阈值，但不是每个值都大于阈值的序列（至少两个值）的一部分。这就是为什么这个索引不会出现在我们的输出中。

另一方面，对于索引，[3, 4, 5]我们有一系列（至少两个）相邻值[5, 4, 6]，其中每个值都高于阈值，这就是它们的索引包含在最终输出中的原因！

---

到目前为止我的代码

我已经用这样的方法来处理这个问题：

(arr > 3).nonzero()

上述命令收集高于阈值的所有项目的索引。但是，我无法确定它们是否连续。我曾想过尝试diff对上述代码片段的结果进行尝试，然后可能会定位其中的结果（也就是说，索引是一个接一个）。这会给我们：

np.diff((arr > 3).nonzero())

但我仍然会在这里遗漏一些东西。

Answer 1

如果将一个布尔数组与一个充满1大小win_size( ) 的窗口进行卷积，那么您将获得一个数组，其中包含项目条件所对应的[1] * win_size值：win_sizewin_size

import numpy as np

def groups(arr, *, threshold, win_size, merge_contiguous=False, flat=False):
    conv = np.convolve((arr >= threshold).astype(int), [1] * win_size, mode="valid")
    indexes_start = np.where(conv == win_size)[0]
    indexes = [np.arange(index, index + win_size) for index in indexes_start]
    
    if flat or merge_contiguous:
        indexes = np.unique(indexes)
        if merge_contiguous:
            indexes = np.split(indexes, np.where(np.diff(indexes) != 1)[0] + 1)
    return indexes


arr = np.array([1, 5, 0, 5, 4, 6, 1, -1, 5, 10])
threshold = 3
win_size = 2

print(groups(arr, threshold=threshold, win_size=win_size))
print(groups(arr, threshold=threshold, win_size=win_size, merge_contiguous=True))
print(groups(arr, threshold=threshold, win_size=win_size, flat=True))

[array([3, 4]), array([4, 5]), array([8, 9])]
[array([3, 4, 5]), array([8, 9])]
[3 4 5 8 9]