Search code examples
pythonpython-3.xiteratoraggregategenerator-expression

How to sum variable-size parts of a collection?


I want to calculate the sum of a collection, for sections of different sizes:

d = (1, 2, 3, 4, 5, 6, 7, 8, 9)
sz = (2, 3, 4)

# here I expect 1+2=3, 3+4+5=12, 6+7+8+9=30

itd = iter(d)
result = tuple( sum(tuple(next(itd) for i in range(s))) for s in sz )

print("result = {}".format(result))

I wonder whether the solution I came up with is the most 'pythonic' (elegant, readable, concise) way to achieve what I want...

In particular, I wonder whether there is a way to get rid of the separate iterator 'itd', and whether it would be easier to work with slices?


Solution

  • I would use itertools.islice since you can directly use the values in sz as the step size at each point:

    >>> from itertools import islice
    >>> it=iter(d)
    >>> [sum(islice(it,s)) for s in sz]
    [3, 12, 30]
    

    Then you can convert that to a tuple if needed.

    The iter is certainly needed in order to step through the tuple at the point where the last slice left off. Otherwise each slice would be d[0:s]