Come sort of code challenge

count_stable_particles.py

from itertools import groupby

def cumulative(arr):
    return [arr[i] - arr[i - 1] for i in range(1, len(arr))]


def stable_particle(arr):
    """
    >>> stable_particle([-1, 1, 3, 3, 3, 2, 3, 2, 1, 0])
    5
    >>> stable_particle([2] * 10000)
    49985001
    >>> stable_particle([1, 3, 5, 7, 9])
    6
    >>> stable_particle([7, 7, 7, 7])
    3
    """
    x = cumulative(arr)
    xx = cumulative(x)
    count_stable = 0
    for k, v in groupby(xx):
        if k == 0:
            n = len(list(v))
            count_stable += (n * (n + 1)) // 2
    return count_stable

def stable_particle2(arr):
    """
    >>> stable_particle2([-1, 1, 3, 3, 3, 2, 3, 2, 1, 0])
    5
    >>> stable_particle2([2] * 10000)
    49985001
    >>> stable_particle2([1, 3, 5, 7, 9])
    6
    >>> stable_particle2([7, 7, 7, 7])
    3
    """
    # Rewrite code with generator
    x = cumulative(arr)
    xx = cumulative(x)
    group_by_gen = (len(list(v)) for k, v in groupby(xx) if k == 0)
    return sum((n * (n + 1)) // 2 for n in group_by_gen)


def main(particles):
    print(stable_particle(particles))


if __name__ == '__main__':
    from doctest import testmod
    testmod()