如何根据谓词分割序列?

Question

我经常需要把一串数据分成两个部分：一部分是符合某个条件的元素，另一部分是不符合的，同时还要保持原来的顺序。

这个假想中的“分割器”函数在实际操作中可能是这样的：

>>> data = map(str, range(14))
>>> pred = lambda i: int(i) % 3 == 2
>>> splitter(data, pred)
[('2', '5', '8', '11'), ('0', '1', '3', '4', '6', '7', '9', '10', '12', '13')]

我想问的是：

Python有没有现成的方法可以做到这一点呢？

其实写这个功能并不难（见下面的附录），但出于一些原因，我更希望能用现成的方法，而不是自己写一个。

谢谢！

附录：

到目前为止，我找到的处理这个任务的最好标准函数是 itertools.groupby。不过，要用它来完成这个特定的任务，必须对每个列表中的元素调用条件函数两次，这让我觉得有点烦人：

>>> import itertools as it
>>> [tuple(v[1]) for v in it.groupby(sorted(data, key=pred), key=pred)]
[('0', '1', '3', '4', '6', '7', '9', '10', '12', '13'), ('2', '5', '8', '11')]

（上面的最后输出和之前想要的结果不同，因为符合条件的元素部分在最后，而不是在最前面，不过这只是个小问题，如果需要的话很容易修正。）

可以通过做一些“内联记忆化”来避免重复调用条件函数，但我尝试的方式变得相当复杂，远没有 splitter(data, pred) 那么简单：

>>> first = lambda t: t[0]
>>> [zip(*i[1])[1] for i in it.groupby(sorted(((pred(x), x) for x in data),
... key=first), key=first)]
[('0', '1', '3', '4', '6', '7', '9', '10', '12', '13'), ('2', '5', '8', '11')]

顺便说一下，如果你不在乎保持原来的顺序，使用 sorted 的默认排序就可以完成这个任务（所以在 sorted 的调用中可以省略 key 参数）：

>>> [zip(*i[1])[1] for i in it.groupby(sorted(((pred(x), x) for x in data)),
... key=first)]
[('0', '1', '3', '4', '6', '7', '9', '10', '12', '13'), ('2', '5', '8', '11')]

Answer 1

分区是一个很实用的功能，属于一些itertools的配方。它使用tee()这个工具，确保在处理集合时，即使有多个迭代器，也只需一次遍历。还用到了内置的filter()函数来筛选出符合条件的项目，以及filterfalse()来获取不符合条件的项目。这是你能找到的最接近标准或内置方法的方式。

def partition(pred, iterable):
    'Use a predicate to partition entries into false entries and true entries'
    # partition(is_odd, range(10)) --> 0 2 4 6 8   and  1 3 5 7 9
    t1, t2 = tee(iterable)
    return filterfalse(pred, t1), filter(pred, t2)

Answer 2

我知道你说你不想写自己的函数，但我真想不通为什么。你的解决方案其实也在写代码，只是没有把它们整理成函数而已。

这个代码正好满足你的需求，容易理解，而且每个元素只会检查一次条件：

def splitter(data, pred):
    yes, no = [], []
    for d in data:
        if pred(d):
            yes.append(d)
        else:
            no.append(d)
    return [yes, no]

如果你想让代码看起来更简洁（不知道为什么会想这样）：

def splitter(data, pred):
    yes, no = [], []
    for d in data:
        (yes if pred(d) else no).append(d)
    return [yes, no]

Answer 3

在 more_itertools 这个库里，有一个叫 partition 的函数，它正好可以完成提问者想要的功能。

from more_itertools import partition

numbers = [1, 2, 3, 4, 5, 6, 7]
predicate = lambda x: x % 2 == 0

predicate_false, predicate_true = partition(predicate, numbers)

print(list(predicate_false), list(predicate_true))

结果是 [1, 3, 5, 7] [2, 4, 6]。