Run Length Encoding in Python

rle.py

"""Run-Length Encoding."""

import numpy as np
import pandas as pd
from itertools import chain


def rle_encode(ndarr):
    """
    Input: ndarray

    Example
    -------

        >>> t, f = True, False
        >>> s = pd.Series([t,t,t,f,t,t,f,f,t,t,t,f,f,f,f,f])
        >>> # encode and decode right back:
        >>> rle_decode(*rle_encode(s.values))
    """
    assert len(ndarr.shape) == 1, "Submit 1-dim arrays only!"
    n = len(ndarr)
    if n == 0:
        return np.array([], dtype=bool), np.array([], dtype=int)
    jump_indices_and_n = np.append(np.nonzero(ndarr[1:] != ndarr[:-1])[0], n-1)
    values = ndarr[jump_indices_and_n]
    lengths = np.diff(np.insert(jump_indices_and_n, 0, -1))

    return values, lengths


def flatmap(f, items):
        return list(chain.from_iterable(map(f, items)))


def rle_decode(values, lengths):
    return np.array(flatmap(lambda vl: [vl[0]] * vl[1], zip(values, lengths)))