"Pythonic" way to return elements from an iterable as long as a condition based on previous element is true

Question

I am working on some code that needs to constantly take elements from an iterable as long as a condition based on (or related to) the previous element is true. For example, let's say I have a list of numbers:

lst = [0.1, 0.4, 0.2, 0.8, 0.7, 1.1, 2.2, 4.1, 4.9, 5.2, 4.3, 3.2]

And let's use a simple condition: the number does not differ from the previous number more than 1. So the expected output would be

[0.1, 0.4, 0.2, 0.8, 0.7, 1.1]

Normally, itertools.takewhile would be a good choice, but in this case it's a bit annoying because the first element doesn't have a previous element to query. The following code returns an empty list because for the first element the code queries the last element.

from itertools import takewhile
res1 = list(takewhile(lambda x: abs(lst[lst.index(x)-1] - x) <= 1., lst))
print(res1)
# []

I managed to write some "ugly" code to work around:

res2 = []
for i, x in enumerate(lst):
    res2.append(x)
    # Make sure index is not out of range
    if i < len(lst) - 1:
        if not abs(lst[i+1] - x) <= 1.:
            break
print(res2)
# [0.1, 0.4, 0.2, 0.8, 0.7, 1.1]

However, I feel like there should be more "pythonic" way to code this. Any suggestions?

Answer 1

You can write your own version of takewhile where the predicate takes both the current and previous values:

def my_takewhile(iterable, predicate):
    iterable = iter(iterable)
    try:
        previous = next(iterable)
    except StopIteration:
        # next(iterable) raises if the iterable is empty
        return
    yield previous
    for current in iterable:
        if not predicate(previous, current):
            break
        yield current
        previous = current

Example:

>>> list(my_takewhile(lst, lambda x, y: abs(x - y) <= 1))
[0.1, 0.4, 0.2, 0.8, 0.7, 1.1]

Answer 2

Solution using assignment expression := for Python >= 3.8:

lst = [0.1, 0.4, 0.2, 0.8, 0.7, 1.1, 2.2, 4.1, 4.9, 5.2, 4.3, 3.2]

pred = lambda cur, prev: abs(cur-prev) <= 1
p = None
res = [p := i for i in lst if p is None or pred(p, i)]

Answer 3

Create a sequence of tuples by zipping the list with a sequence that prepends the first element of the list to the list; the resulting sequence of tuples pairs the first element with itself (so that abs(x-x) is guaranteed less than 1) and each other element with its preceding element.

a = lst                        == x1       x2       x3       x4       ...
b = chain(islice(lst, 1), lst) == x1       x1       x2       x3       ...
zip(a, b)                      == (x1, x1) (x2, x1) (x3, x2) (x4, x3) ...

Then

>>> from itertools import takewhile, chain
>>> lst = [0.1, 0.4, 0.2, 0.8, 0.7, 1.1, 2.2, 4.1, 4.9, 5.2, 4.3, 3.2]
>>> def close(t): return abs(t[0] - t[1]) <= 1
...
>>> [x for x, _ in takewhile(close, zip(lst, chain(islice(lst, 1), lst)))]
[0.1, 0.4, 0.2, 0.8, 0.7, 1.1]

If you prefer, you can define prepend as shown in the itertools documentation, and write

[x for x, _ in takewhile(close, zip(lst, prepend(lst[0], lst)))]

You can also just use ordinary list slicing in this case instead of islice (which is essentially just inlining the aforementioned prepend function, as lst[:1] == [lst[0]]).

[x for x, _ in takewhile(close, zip(lst, chain(lst[:1], lst)))]