如何在谓词首次为False处将列表分为两部分

Question

我一直在想应该有个函数可以做到这一点，但我在可能的地方（谷歌、itertools文档、列表方法、其他StackOverflow问题）搜索过，却没有找到我想要的东西。

这是一个简单但有效的实现：

def split_at_first_false(pred, seq):
    first = []
    second = []
    true_so_far = True
    for item in seq:
        if true_so_far and pred(item):
            first.append(item)
        else:
            true_so_far = False
            second.append(item)
    return first, second

print split_at_first_false(str.isalpha, "abc1a2b")
# (['a', 'b', 'c'], ['1', 'a', '2', 'b'])

它能工作，但感觉不太对。应该有更好的方法来做这件事！

编辑：在查看了答案后，我最终使用了senderle最后建议的一个稍微修改过的版本：

from itertools import chain

def split_at_pred(pred, seq):
    head = []
    it = iter(seq)
    for i in it:
        if not pred(i):
            head.append(i)
        else:
            return iter(head), chain([i], it)
    return iter(head), iter([])

这个方法简短而优雅，无论输入是什么（字符串、列表、迭代器），输出都是两个迭代器，作为额外的好处，它甚至可以处理以下输入：

from itertools import count
split_at_pred(lambda x: x == 5, count())

其他解决方案，即使能处理迭代器，也会在这个输入下耗尽内存。（注意，这个确实只是一个额外的好处。在我写这个问题时，我甚至没有考虑到无限迭代器）

Answer 1

不要害怕使用迭代器，这正是使用它的好时机。一旦遇到第一个不符合条件的项目，就可以用同一个迭代器把剩下的项目填入第二个列表。

def split_at_false(pred, seq):
    # if seq is not already an iterator, make it one
    if not hasattr(seq,'next'):
        seq = iter(seq)

    first, second = [], []
    for item in seq:
        if not pred(item):
            second.append(item)
            break
        first.append(item)

    # at this point, seq points to the first item
    # after the false item, just add it and all the 
    # rest to the second list
    second.extend(seq)

    return first, second

is_odd = lambda x : x % 2    
print split_at_false(is_odd, [1])    
print split_at_false(is_odd, [1,2,3,4,5])
print split_at_false(is_odd, [2,3,4,5,6])
print split_at_false(is_odd, [])

输出结果：

([1], [])
([1], [2, 3, 4, 5])
([], [2, 3, 4, 5, 6])
([], [])

没有额外的列表存储，没有重复遍历列表，也没有切片操作，就是一个简单的迭代器。

Answer 2

这样怎么样？

def split_at_first_false(pred, seq):
    for i, item in enumerate(seq):
        if not pred(item):
            return seq[:i], seq[i:]

Answer 3

这看起来是个适合用itertools来解决的问题。

>>> first = list(itertools.takewhile(str.isalpha, l))
>>> second = list(itertools.dropwhile(str.isalpha, l))
>>> first
['a', 'b', 'c']
>>> second
['1', 'a', '2', 'b']

如果l是一个迭代器而不是一个序列，那么这个部分需要做一些调整。

>>> def bisect_iter(pred, i):
...     i1, i2 = itertools.tee(i)
...     return itertools.takewhile(pred, i1), itertools.dropwhile(pred, i2)
... 
>>> i1, i2 = bisect_iter(str.isalpha, iter(l))
>>> list(i1)
['a', 'b', 'c']
>>> list(i2)
['1', 'a', '2', 'b']

使用tee的缺点是，初始值会被缓存，并且会被测试两次（分别由takewhile和dropwhile进行）。这样做有点浪费。不过，如果你想同时接受和返回迭代器，缓存值是不可避免的。

不过，如果你能从一个迭代器返回列表，我想到一个解决方案，它不会额外复制或测试，而且和你的方法非常接近：

>>> def bisect_iter_to_list(pred, it):
...     l1 = []
...     for i in it:
...         if pred(i):
...             l1.append(i)
...         else:
...             l2 = [i]
...             l2.extend(it)
...     return l1, l2
... 
>>> bisect_iter_to_list(str.isalpha, iter(l))
(['a', 'b', 'c'], ['1', 'a', '2', 'b'])

唯一需要注意的是，通常在else语句后面会有一个break语句，但我只是消费了迭代器，这样for循环就会提前结束。

最后，如果你仍然想返回迭代器，但又不想做额外的测试，这里有一个我认为是最优的变体。

>>> def bisect_any_to_iter(pred, it):
...     it = iter(it)
...     head = []
...     for i in it:
...         if pred(i):
...             head.append(i)
...         else:
...             tail = itertools.chain([i], it)
...             break
...     return iter(head), tail
... 
>>> a, b = bisect_iter_to_iter(str.isalpha, iter(l))
>>> list(a)
['a', 'b', 'c']
>>> list(b)
['1', 'a', '2', 'b']