nth_iterinterleave.py

def nth_iterinterleave(*arrays, nth=1):
    """Interleave multiple lists.
       via https://github.com/dgilland/pydash/blob/develop/src/pydash/arrays.py
       
      Expanded to allow an nth parameter, where arrays[0] is treated as the primary and arrays[1:] are interleaved every n elements

      Example:

      given *arrays = [ [1,2,3,4,5,6,7,8,9,10], ['x', 'y', 'z'], ['a', 'b', 'c'] ]

      nth=1 (default) returns:
      [1, 'x', 'a', 2, 'y', 'b', '3', 'z', 'c', 4, 5, 6, 7, 8, 9, 10]

      nth=2 returns:
      [1, 2, 'x', 'a', 3, 4, 'y', 'b', 5, 6, 'z', 'c', 7, 8, 9, 10]

      nth=3 returns:
      [1, 2, 3, 'x', 'a', 4, 5, 6, 'y', 'b', 7, 8, 9, 'z', 'c', 10]

      via https://github.com/dgilland/pydash/blob/develop/src/pydash/arrays.py
    """

    iters = [iter(arr) for arr in arrays]

    # We'll increment a counter each time we yield an element from the primary array
    # When counter % nth equals 0, we know it's time to interleave elements from the non-primary arrays
    counter = 1
    while iters:
        nextiters = []
        for i, itr in enumerate(iters):
            # primary array
            # always yield next item
            if i == 0:
                counter += 1
                try:
                    yield next(itr)
                    nextiters.append(itr)
                except StopIteration:
                    pass
            # secondary arrays
            # only yield next items if counter % nth == 0
            else:
                if counter % nth == 0:
                    try:
                        counter = 1
                        yield next(itr)
                        nextiters.append(itr)
                    except StopIteration:
                        pass
                else:
                    nextiters.append(itr)

        iters = nextiters