Count lengths of constant subarrays

Question

I have an array assuming discrete values. For each value, I am interested in how often it is assumed in a row.
For instance, if my array is [5 5 5 5 5 0 0 4 1 1 5 1 1 1 1 1 3 3 3 3 2 2 5 1 1 1 1 1 1 2 2 5], then I expect as a result the following dictionary:

{
    0: [2],
    1: [2, 5, 6],
    2: [2, 2],
    3: [4],
    4: [1],
    5: [5, 1, 1, 1]
}

The keys are the discrete values from the array. The value for each key is a list telling me how often the key was hit in a row when it occured.

My (working) approach is this:

def durations_per_occurence(
    array: np.ndarray,
    max: int,     # max == np.max(array)
) -> dict:
    durations = {k: [] for k in range(max+1)}
    
    current_key = array[0]
    current_duration = 0
    for key in array:
        if key == current_key:
            current_duration += 1
        else:
            durations[current_key].append(current_duration)
            current_key = key
            current_duration = 1
    
    durations[current_key].append(current_duration) # last entry
    
    return durations

I am grateful for any advice, but particulary interested in built-in Python functionality to speed up the computation. (I often hear that for-loops are particularly slow in Python and that one should vectorize computations. Is there a neat vectorization for this particular problem?)

Answer 1

A few initial thoughts on the code:

• Iterating through groups of consecutive elements is implemented efficiently by itertools.groupby, which for each consecutive group in a collection returns both the repeated value and an iterator through those repeated elements in the list.
• cardinality.count can efficiently count the number of elements in that iterable, enabling us to efficiently get the run length
• In cases where you have a vector with a wide range of data values, you may be creating a durations dict that is mostly empty, spending time and memory in the process. Plus you need to compute the maximum value, which takes time, adds an extra function argument, and is prone to off-by-one errors. It also reduces flexibility if, for instance, you started getting negative or fractional data (which break your code). The collections.defaultdict object smoothly handles all these issues -- it automagically creates the initial empty list for each key when it's first added, and it doesn't store keys that never show up (but it will return [] if an unadded key is accessed).

Putting it all together, you end up with:

from itertools import groupby
from cardinality import count
from collections import defaultdict

def durations_per_occurence2(
    array: np.ndarray
) -> defaultdict:
    durations = defaultdict(list)
    for k, g in groupby(array):
        durations[k].append(count(g))
    return durations

Beyond being much more compact and pythonic, this code runs more than 10,000x faster (according to timeit) when I run with your sample data with the number 1000000 appended to the end.