Iterate over a ‘window’ of adjacent elements in Python

Viewed 3837

This is more a question of elegance and performance rather than “how to do at all”, so I'll just show the code:

def iterate_adjacencies(gen, fill=0, size=2, do_fill_left=True,
  do_fill_right=False):
    """ Iterates over a 'window' of `size` adjacent elements in the supploed
    `gen` generator, using `fill` to fill edge if `do_fill_left` is True
    (default), and fill the right edge (i.e.  last element and `size-1` of
    `fill` elements as the last item) if `do_fill_right` is True.  """
    fill_size = size - 1
    prev = [fill] * fill_size
    i = 1
    for item in gen:  # iterate over the supplied `whatever`.
        if not do_fill_left and i < size:
            i += 1
        else:
            yield prev + [item]
        prev = prev[1:] + [item]
    if do_fill_right:
        for i in range(fill_size):
            yield prev + [fill]
            prev = prev[1:] + [fill]

and then ask: is there already a function for that? And, if not, can you do the same thing in a better (i.e. more neat and/or more fast) way?

Edit:

with ideas from answers of @agf, @FogleBird, @senderle, a resulting somewhat-neat-looking piece of code is:

def window(seq, size=2, fill=0, fill_left=True, fill_right=False):
    """ Returns a sliding window (of width n) over data from the iterable:
      s -> (s0,s1,...s[n-1]), (s1,s2,...,sn), ...
    """
    ssize = size - 1
    it = chain(
      repeat(fill, ssize * fill_left),
      iter(seq),
      repeat(fill, ssize * fill_right))
    result = tuple(islice(it, size))
    if len(result) == size:  # `<=` if okay to return seq if len(seq) < size
        yield result
    for elem in it:
        result = result[1:] + (elem,)
        yield result
5 Answers
Related