如何在NumPy数组中找到连续元素组

91 投票
7 回答
64338 浏览
提问于 2025-04-17 01:49

我需要把NumPy数组中的连续元素分组。下面是一个例子:

a = [ 0, 47, 48, 49, 50, 97, 98, 99]

输出应该是一个包含元组的列表,像这样:

[(0), (47, 48, 49, 50), (97, 98, 99)]

这里元素之间的差值只有1。如果能把这个差值设定成一个限制值或者一个固定的数字,那就更好了。

7 个回答

13

(a[1:]-a[:-1])==1 这段代码会生成一个布尔数组,其中 False 表示在连续的数字中出现了断裂。你也可以使用内置的 numpy.grad 函数。

238
def consecutive(data, stepsize=1):
    return np.split(data, np.where(np.diff(data) != stepsize)[0]+1)

a = np.array([0, 47, 48, 49, 50, 97, 98, 99])
consecutive(a)
[array([0]), array([47, 48, 49, 50]), array([97, 98, 99])]

产生

25

这里有一个小函数,可能会对你有帮助:

def group_consecutives(vals, step=1):
    """Return list of consecutive lists of numbers from vals (number list)."""
    run = []
    result = [run]
    expect = None
    for v in vals:
        if (v == expect) or (expect is None):
            run.append(v)
        else:
            run = [v]
            result.append(run)
        expect = v + step
    return result

>>> group_consecutives(a)
[[0], [47, 48, 49, 50], [97, 98, 99]]
>>> group_consecutives(a, step=47)
[[0, 47], [48], [49], [50, 97], [98], [99]]

附注:这段代码是纯Python写的。如果你想要NumPy的解决方案,可以看看unutbu的回答。

撰写回答