From 45ff04ac75ddeb70fe4647a146790d17c3f3b8f8 Mon Sep 17 00:00:00 2001 From: stijn Date: Mon, 24 Apr 2017 15:56:44 +0200 Subject: [PATCH] itertools: Add accumulate function --- itertools/itertools.py | 11 +++++++++++ itertools/test_itertools.py | 6 ++++++ 2 files changed, 17 insertions(+) diff --git a/itertools/itertools.py b/itertools/itertools.py index 5c00952f..2ff05347 100644 --- a/itertools/itertools.py +++ b/itertools/itertools.py @@ -55,3 +55,14 @@ def tee(iterable, n=2): def starmap(function, iterable): for args in iterable: yield function(*args) + +def accumulate(iterable, func=lambda x, y: x + y): + it = iter(iterable) + try: + acc = next(it) + except StopIteration: + return + yield acc + for element in it: + acc = func(acc, element) + yield acc diff --git a/itertools/test_itertools.py b/itertools/test_itertools.py index 6a4ad1be..1481e8ac 100644 --- a/itertools/test_itertools.py +++ b/itertools/test_itertools.py @@ -14,3 +14,9 @@ assert list(itertools.islice(itertools.cycle([1, 2, 3]), 10)) == [1, 2, 3, 1, 2, assert list(itertools.islice(itertools.cycle(reversed([1, 2, 3])), 7)) == [3, 2, 1, 3, 2, 1, 3] assert list(itertools.starmap(lambda x, y: x * y, [[1, 2], [2, 3], [3, 4]])) == [2, 6, 12] + +assert list(itertools.accumulate([])) == [] +assert list(itertools.accumulate([0])) == [0] +assert list(itertools.accumulate([0, 2, 3])) == [0, 2, 5] +assert list(itertools.accumulate(reversed([0, 2, 3]))) == [3, 5, 5] +assert list(itertools.accumulate([1, 2, 3], lambda x, y: x * y)) == [1, 2, 6]