How to elegantly generate all prefixes of an iterable? (cumulative iterable)

691 Views Asked by At

From an iterable, I'd like to generate an iterable of its prefixes (including the original iterable itself).

for prefix in prefixes(range(5)):
    print(tuple(prefix))

should result in

(0,)
(0, 1)
(0, 1, 2)
(0, 1, 2, 3)
(0, 1, 2, 3, 4)

or in

()
(0,)
(0, 1)
(0, 1, 2)
(0, 1, 2, 3)
(0, 1, 2, 3, 4)

and

for prefix in prefixes('Hello'):
    print(''.join(prefix))

should result in

H
He
Hel
Hell
Hello

or in


H
He
Hel
Hell
Hello

(Whether the empty prefix is part of the result doesn't matter too much for me, nor does the exact type of the inner or outer resulting iterables.)

I was able to devise several ways to implement this, but all feel at least slightly clunky:

using slicing & len:

(works if the iterable is a sequence)

def prefixes(seq):
    for i in range(len(seq)):
        yield seq[:i + 1]

or using a list comprehension:

def prefixes(seq):
    return [seq[:i + 1] for i in range(len(seq))]

... or a generator expression

def prefixes(seq):
    return (seq[:i + 1] for i in range(len(seq)))

(These don't yield the empty prefix. To include it, replace [i + 1] by just [i] and range(len(seq)) by range(len(seq) + 1) in any of the above.)

These feel clunky:

  • because they don't work for all kinds iterable inputs
  • because of the need for the + 1 offset
  • calling range on the len of something (though enumerate wouldn't make it better here)

using concatenation

def prefixes(iterable):
    result = ()
    for elem in iterable:
        result += (elem,)
        yield result

(Doesn't include the empty prefix. This can be changed by yielding result already once before the for-loop.)

or using itertools.accumulate

from itertools import accumulate as acc

def prefixes(iterable):
    return acc(iterable, lambda t, elem: t + (elem,), initial=())

or a bit more readable:

from itertools import accumulate

def _append(iterable, elem):
    return iterable + (elem,)

def prefixes(iterable):
    return accumulate(iterable, _append, initial=())

(These two include the empty prefix. Drop it if unwanted.)

These feel clunky due to the need to pack elements into length-one containers just to concatenate them to an existing one.

Solutions that are more elegant?

I feel like I must be missing something from itertools, functools, operator or more-itertools that would allow for a slightly or even significantly less clunky implementation. I mean, this is eerily similar to more_itertools.powerset, just a, well, rather specific subset of it.

3

There are 3 best solutions below

0
On BEST ANSWER

It may be considered elegant to write the prefixes function in any generalized way that works, put it in a module, and then import it in the code where it is needed, so that it doesn't matter how it is implemented.

On the other hand, requiring an extra import can be perceived as less elegant than a short local function that is less generic but more tailored to the specific use case.

This is one possible quite generic solution:

def prefixes(iterable):
    return itertools.accumulate(map(lambda x: (x,), iterable))

There are reasons for it to be considered elegant:

  • It uses a function that is already available in the standard library and achieves the primary goal,
  • it does not explicitly mention the concatenation which accumulate already does implicitly,
  • it does not require the initial argument to accumulate.

But some find using map and lambda to be less elegant than a for loop.

2
On

This isn't fully fleshed-out, and it's also a bit dorky:

def prefixes(iterable):
    from itertools import tee, islice
    iterator = iter(iterable)
    length = len(iterable)
    for slice_length, it in enumerate(tee(iterator, length), start=1):
        yield islice(it, slice_length)


for prefix in prefixes(range(5)):
    print(tuple(prefix))

for prefix in prefixes("Hello"):
    print("".join(prefix))

Output:

(0,)
(0, 1)
(0, 1, 2)
(0, 1, 2, 3)
(0, 1, 2, 3, 4)
H
He
Hel
Hell
Hello

You end up making n+1 independent iterators of the iterable. You also need to know the length of the iterable in advance, or be able to take the length of it (so you can't pass in a generator.)

4
On

Similar to your first concatenation example, but building a list instead of a tuple:

def prefixes(iterable):
    result = []
    for elem in iterable:
        result.append(elem)
        yield result

This eliminates the necessity of creating temporary one-element tuples.