"Pythonic" 从迭代中获取 return 元素的方法,只要基于前一个元素的条件为真

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

我正在编写一些代码,只要基于(或相关)前一个元素的条件为真,就需要不断地从可迭代对象中获取元素。例如,假设我有一个数字列表:

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 会是一个不错的选择,但在这种情况下有点烦人,因为第一个元素没有要查询的前一个元素。以下代码 returns 一个空列表,因为对于第一个元素,代码查询最后一个元素。

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”方式来编写代码。有什么建议吗?

您可以编写自己的 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]

我会推荐一个简单的 for 循环。拥有一个可以操作的整数将使您可以轻松地比较同一列表中的多个值。

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

索引从 1 开始而不是零,以便您可以比较初始的两个值。

或者,如果您需要对第一个元素做一些特殊的事情,那么从 0 开始 for 循环并为第一种情况添加一个简单的 if 语句:

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

通过使用将列表的第一个元素添加到列表中的序列压缩列表来创建元组序列;结果序列 元组的第一个元素与其自身配对(因此 abs(x-x) 保证小于 1),每个其他元素与其前面的元素配对。

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

然后

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

如果您愿意,可以按照 itertools 文档中所示定义 prepend,然后编写

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

在这种情况下,您也可以只使用普通列表切片而不是 islice(本质上只是内联上述 prepend 函数,如 lst[:1] == [lst[0]])。

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

解决方案使用赋值表达式 := 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)]

稍后但是,这里有另一种解决方案,也许不是最Pythonic方式:

也许你可以考虑递归的方法。函数 reduce_list 接收一个列表(您的 lst 变量)和该列表的当前项(第一个)作为参数。有一个名为 list_result 的变量将存储满足条件的项目。如果列表只有一项,那么我们就没有什么可以比较它的了,因此我们将 return list_result。否则,我们获取列表中的下一项(第二项),如果条件为 True,则我们存储该项目。如果条件是False,那么我们保存最后一个current并停止递归。

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)

输出:[0.1, 0.4, 0.2, 0.8, 0.7, 1.1]