只要基于前一个元素的条件为真，就可以从可迭代对象中返回元素的“Pythonic”方式

Question

我正在处理一些需要不断从迭代中获取元素的代码，只要基于（或与之相关）前一个元素的条件为真。例如，假设我有一个数字列表：

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]

让我们使用一个简单的条件：该数字与前一个数字的差异不超过 1。因此预期输出将是

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

通常，itertools.takewhile这是一个不错的选择，但在这种情况下它有点烦人，因为第一个元素没有要查询的前一个元素。以下代码返回一个空列表，因为对于第一个元素，代码会查询最后一个元素。

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

我设法编写了一些“丑陋”的代码来解决：

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]

但是，我觉得应该有更多的“pythonic”方式来对此进行编码。有什么建议？

Answer 1

您可以编写自己的版本，takewhile其中谓词采用当前值和先前值：

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

例子：

>>> 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

使用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)]