"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)))]

Answer 4

A little later but, here there is another solution, perhaps is not the most Pythonic way:

Maybe you could consider a recursive approach. The function reduce_list receive as parameter a list (your lst variable) and the current item of this list (the first one). There is a variable named list_result that will store the items that accomplish the condition. If the list has only one item, then we don't have against what compare it, for that reason we will return list_result. Otherwise We get the next item in the list (the second one) and if the condition is True, then We store that item. If the condition is False, then We save the last one current and stop the recursion.

list_result = []
def reduce_list(lst, current):
  if len(lst) == 1:
    return list_result
  
  the_next = lst[1]
  if (abs(current - the_next) <= 1):
    list_result.append(current)
  else:
    list_result.append(current)
    return
  
  lst = lst[1:] # Set the lst as the rest of the list itself.
  current = lst[0] # current will be the first one of the rest.
  reduce_list(lst, current)

reduce_list(lst, lst[0])
print(list_result)

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

Answer 5

I would recommend a simple for loop. Having an integer you can manipulate will allow you to easily compare multiple values within the same list.

for i in range(1, len(lst)):
  if(abs(lst[i] - lst[i-1]) < 1):
    # do stuff

The index starts at 1 and not zero so that you can compare the initial two values.

Alternatively, if you need to do something special with the first element, then start the for loop at 0 and add a simple if statement for the first case:

  if(i == 0):
    # do stuff
  elif( normal condition ):