我有这样的代码:

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

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

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


当前回答

下面是惰性迭代器方法:

from itertools import tee

def split_on_condition(seq, condition):
    l1, l2 = tee((condition(item), item) for item in seq)
    return (i for p, i in l1 if p), (i for p, i in l2 if not p)

它对每个项计算一次条件,并返回两个生成器,第一个生成条件为真时序列中的值,另一个生成条件为假时序列中的值。

因为它是惰性的,你可以在任何迭代器上使用它,甚至是无限迭代器:

from itertools import count, islice

def is_prime(n):
    return n > 1 and all(n % i for i in xrange(2, n))

primes, not_primes = split_on_condition(count(), is_prime)
print("First 10 primes", list(islice(primes, 10)))
print("First 10 non-primes", list(islice(not_primes, 10)))

通常情况下,非惰性列表返回方法会更好:

def split_on_condition(seq, condition):
    a, b = [], []
    for item in seq:
        (a if condition(item) else b).append(item)
    return a, b

编辑:对于您更具体的用例,将项目按某些键分割到不同的列表中,这里有一个通用函数:

DROP_VALUE = lambda _:_
def split_by_key(seq, resultmapping, keyfunc, default=DROP_VALUE):
    """Split a sequence into lists based on a key function.

        seq - input sequence
        resultmapping - a dictionary that maps from target lists to keys that go to that list
        keyfunc - function to calculate the key of an input value
        default - the target where items that don't have a corresponding key go, by default they are dropped
    """
    result_lists = dict((key, []) for key in resultmapping)
    appenders = dict((key, result_lists[target].append) for target, keys in resultmapping.items() for key in keys)

    if default is not DROP_VALUE:
        result_lists.setdefault(default, [])
        default_action = result_lists[default].append
    else:
        default_action = DROP_VALUE

    for item in seq:
        appenders.get(keyfunc(item), default_action)(item)

    return result_lists

用法:

def file_extension(f):
    return f[2].lower()

split_files = split_by_key(files, {'images': IMAGE_TYPES}, keyfunc=file_extension, default='anims')
print split_files['images']
print split_files['anims']

其他回答

第一步(pre-OP-edit):使用集合:

mylist = [1,2,3,4,5,6,7]
goodvals = [1,3,7,8,9]

myset = set(mylist)
goodset = set(goodvals)

print list(myset.intersection(goodset))  # [1, 3, 7]
print list(myset.difference(goodset))    # [2, 4, 5, 6]

这对可读性(IMHO)和性能都有好处。

第二步(post-OP-edit):

创建一个好的扩展列表:

IMAGE_TYPES = set(['.jpg','.jpeg','.gif','.bmp','.png'])

这将提高性能。否则,你现在的情况在我看来还不错。

之前的答案似乎并不能满足我所有的四种强迫症:

尽可能的懒惰, 只对原始Iterable求值一次 每个项只计算谓词一次 提供良好的类型注释(适用于python 3.7)

我的解决方案并不漂亮,我不认为我可以推荐使用它,但它是:

def iter_split_on_predicate(predicate: Callable[[T], bool], iterable: Iterable[T]) -> Tuple[Iterator[T], Iterator[T]]:
    deque_predicate_true = deque()
    deque_predicate_false = deque()
    
    # define a generator function to consume the input iterable
    # the Predicate is evaluated once per item, added to the appropriate deque, and the predicate result it yielded 
    def shared_generator(definitely_an_iterator):
        for item in definitely_an_iterator:
            print("Evaluate predicate.")
            if predicate(item):
                deque_predicate_true.appendleft(item)
                yield True
            else:
                deque_predicate_false.appendleft(item)
                yield False
    
    # consume input iterable only once,
    # converting to an iterator with the iter() function if necessary. Probably this conversion is unnecessary
    shared_gen = shared_generator(
        iterable if isinstance(iterable, collections.abc.Iterator) else iter(iterable)
    )
    
    # define a generator function for each predicate outcome and queue
    def iter_for(predicate_value, hold_queue):
        def consume_shared_generator_until_hold_queue_contains_something():
            if not hold_queue:
                try:
                    while next(shared_gen) != predicate_value:
                        pass
                except:
                    pass
        
        consume_shared_generator_until_hold_queue_contains_something()
        while hold_queue:
            print("Yield where predicate is "+str(predicate_value))
            yield hold_queue.pop()
            consume_shared_generator_until_hold_queue_contains_something()
    
    # return a tuple of two generators  
    return iter_for(predicate_value=True, hold_queue=deque_predicate_true), iter_for(predicate_value=False, hold_queue=deque_predicate_false)

用下面的测试,我们从print语句中得到如下输出:

t,f = iter_split_on_predicate(lambda item:item>=10,[1,2,3,10,11,12,4,5,6,13,14,15])
print(list(zip(t,f)))
# Evaluate predicate.
# Evaluate predicate.
# Evaluate predicate.
# Evaluate predicate.
# Yield where predicate is True
# Yield where predicate is False
# Evaluate predicate.
# Yield where predicate is True
# Yield where predicate is False
# Evaluate predicate.
# Yield where predicate is True
# Yield where predicate is False
# Evaluate predicate.
# Evaluate predicate.
# Evaluate predicate.
# Evaluate predicate.
# Yield where predicate is True
# Yield where predicate is False
# Evaluate predicate.
# Yield where predicate is True
# Yield where predicate is False
# Evaluate predicate.
# Yield where predicate is True
# Yield where predicate is False
# [(10, 1), (11, 2), (12, 3), (13, 4), (14, 5), (15, 6)]
images = [f for f in files if f[2].lower() in IMAGE_TYPES]
anims  = [f for f in files if f not in images]

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

简单的生成器版本,在内存中保存尽可能少的值,并且只调用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()

这是这个问题的另一个解决方案。我需要一个尽可能快的解决方案。这意味着只对列表进行一次迭代,并且最好是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))。转换要快得多,直接构建列表。

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