CODEWITHSUNDEEP
pythonpython-3.xiteratorpython-itertools
kupgov
kupgov

Reputation: 393

Get pairwise iterator with additional item in the end

The goal: e.g. given finite iterator p0, p1, ..., pn turn into (p0, p1), (p1, p2), ..., (pn-1, pn), (pn, None) — iterator through pairs of consecutive items with special last item.

pairwise() function exists in the documentation as example of itertools usage:

def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

But I want additionally add yet another item to the end of iterator (if it is finite) with some default value for the second element of pair (e.g., None).

How to efficiently implement this additional functionality?

Upvotes: 5

Views: 978

Answers (3)

TigerhawkT3
TigerhawkT3

Reputation: 49318

You can create a generator:

def pairwise(iterable, additional=None):
    iterable = iter(iterable)
    first, second = next(iterable), next(iterable)
    while 1:
        yield first,second
        try:
            first,second = second, next(iterable)
        except TypeError:
            yield second, additional
            break

Results:

>>> list(pairwise([1,2,3], 'a'))
[(1, 2), (2, 3), (3, 'a')]
>>> list(pairwise('abc', 'a'))
[('a', 'b'), ('b', 'c'), ('c', 'a')]
>>> list(pairwise('abcd', 'a'))
[('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'a')]

For an infinite iterable:

>>> a = pairwise(infi(), 6)
>>> for i in range(10):
...     print(next(a))
...
(0, 1)
(1, 2)
(2, 3)
(3, 0)
(0, 1)
(1, 2)
(2, 3)
(3, 0)
(0, 1)
(1, 2)

Upvotes: 0

poke
poke

Reputation: 387745

As for adding (sn, None) at the end, as user2357112 already answered, you can just use zip_longest so one already exhausted iterator is not stopping the whole sequence (so the a iterator can still yield the last element).

For all other situations, e.g. if you want to add further elements at the end, you can just make a generator function itself. All the itertools function are already lazy generators, only producing new results when you request the next element in the result, and you can easily consume those from within a generator.

Let’s say, you need pairwise to yield a sentinel value (None, None) at the end, then you could simply yield the results from zip_longest and then yield another item:

def example (iterable):
    a, b = tee(iterable)
    next(b, None)
    yield from zip_longest(a, b)
    yield (None, None)

The yield from syntax actually came with Python 3.3. For earlier versions, especially Python 2, you would need to do that manually by looping over the items and yielding them again:

def example (iterable):
    a, b = tee(iterable)
    next(b, None)
    for x in zip_longest(a, b):
        yield x
    yield (None, None)

Upvotes: 2

user2357112
user2357112

Reputation: 280963

Using itertools.zip_longest:

def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip_longest(a, b)

When one of the input iterators runs out, zip_longest pads it with a filler value, which defaults to None.

Upvotes: 10

Related Questions