我有这样的代码:

good = [x for x in mylist if x in goodvals]
bad = [x for x in mylist if x not in goodvals]

目标是根据mylist的内容是否满足条件,将它们拆分为另外两个列表。

我怎样才能做得更优雅呢?我能避免在mylist上做两个单独的迭代吗?我可以通过这样做来提高性能吗?


当前回答

简单的生成器版本,在内存中保存尽可能少的值,并且只调用pred一次:

from collections import deque
from typing import Callable, TypeVar, Iterable
_T = TypeVar('_T')

def iter_split(pred: Callable[[_T], bool],
               iterable: Iterable[_T]) -> tuple[Iterable[_T], Iterable[_T]]:
    """Split an iterable into two iterables based on a predicate.
    
    The predicate will only be called once per element.
    
    Returns:
        A tuple of two iterables, the first containing all elements for which
        the predicate returned True, the second containing all elements for
        which the predicate returned False.
    """
    iterator = iter(iterable)
    true_values: deque[_T] = deque()
    false_values: deque[_T] = deque()
    
    def true_generator():
        while True:
            while true_values:
                yield true_values.popleft()
            
            for item in iterator:
                if pred(item):
                    yield item
                    break
                false_values.append(item)
            else:
                break
            
    def false_generator():
        while True:
            while false_values:
                yield false_values.popleft()
            
            for item in iterator:
                if not pred(item):
                    yield item
                    break
                true_values.append(item)
            else:
                break

    return true_generator(), false_generator()

其他回答

受到@gnibbler伟大(但简洁!)回答的启发,我们可以应用该方法映射到多个分区:

from collections import defaultdict

def splitter(l, mapper):
    """Split an iterable into multiple partitions generated by a callable mapper."""

    results = defaultdict(list)

    for x in l:
        results[mapper(x)] += [x]

    return results

然后可以使用splitter,如下所示:

>>> l = [1, 2, 3, 4, 2, 3, 4, 5, 6, 4, 3, 2, 3]
>>> split = splitter(l, lambda x: x % 2 == 0)  # partition l into odds and evens
>>> split.items()
>>> [(False, [1, 3, 3, 5, 3, 3]), (True, [2, 4, 2, 4, 6, 4, 2])]

这适用于有更复杂映射的两个以上分区(也适用于迭代器):

>>> import math
>>> l = xrange(1, 23)
>>> split = splitter(l, lambda x: int(math.log10(x) * 5))
>>> split.items()
[(0, [1]),
 (1, [2]),
 (2, [3]),
 (3, [4, 5, 6]),
 (4, [7, 8, 9]),
 (5, [10, 11, 12, 13, 14, 15]),
 (6, [16, 17, 18, 19, 20, 21, 22])]

或者用字典来映射:

>>> map = {'A': 1, 'X': 2, 'B': 3, 'Y': 1, 'C': 2, 'Z': 3}
>>> l = ['A', 'B', 'C', 'C', 'X', 'Y', 'Z', 'A', 'Z']
>>> split = splitter(l, map.get)
>>> split.items()
(1, ['A', 'Y', 'A']), (2, ['C', 'C', 'X']), (3, ['B', 'Z', 'Z'])]

如果你不想用两行代码来完成一个语义只需要一次的操作,你可以把上面的一些方法(甚至是你自己的方法)包装在一个函数中:

def part_with_predicate(l, pred):
    return [i for i in l if pred(i)], [i for i in l if not pred(i)]

这不是一种惰性计算方法,它确实对列表进行了两次迭代,但是它允许您在一行代码中对列表进行分区。

images = [f for f in files if f[2].lower() in IMAGE_TYPES]
anims  = [f for f in files if f not in images]

当条件较长时很好,例如在您的示例中。读者不需要弄清楚否定条件以及它是否适用于所有其他情况。

为了提高性能,请尝试itertools。

itertools模块标准化了一组快速、内存高效的核心工具,这些工具单独使用或组合使用都很有用。它们一起构成了一个“迭代器代数”,使得用纯Python简洁有效地构造专门的工具成为可能。

出现看到itertools。过滤器或imap。

itertools。iterable ifilter(谓词) 创建一个迭代器,从iterable中过滤元素,只返回谓词为True的元素

这是这个问题的另一个解决方案。我需要一个尽可能快的解决方案。这意味着只对列表进行一次迭代,并且最好是O(1)用于向结果列表之一添加数据。这与sastanin提供的解决方案非常相似,只是更短:

from collections import deque

def split(iterable, function):
    dq_true = deque()
    dq_false = deque()

    # deque - the fastest way to consume an iterator and append items
    deque((
      (dq_true if function(item) else dq_false).append(item) for item in iterable
    ), maxlen=0)

    return dq_true, dq_false

此时,可以按照如下方式使用该函数:

lower, higher = split([0,1,2,3,4,5,6,7,8,9], lambda x: x < 5)

selected, other = split([0,1,2,3,4,5,6,7,8,9], lambda x: x in {0,4,9})

如果你对结果的deque对象不满意,你可以很容易地将其转换为list、set或任何你喜欢的对象(例如list(lower))。转换要快得多,直接构建列表。

该方法保持项目的顺序,以及任何副本。