Python - группировка последовательных элементов, удовлетворяющих предикату

Question

Я хотел бы сгруппировать последовательные элементы, удовлетворяющие предикату. Doctest для такой функции будет выглядеть как

>>> group([1, 2, 3, 0, 4, 5, 0, 0, 6], lambda x: x != 0) 
    [[1, 2, 3], [4, 5], [6]] 
>>> group([1, 2, 3, 0, 4, 5, 0, 0, 6], lambda x: x == 0) 
    [[0], [0, 0]] 

я написал прототип, где я использую itertool в TakeWhile, но это некрасиво, потому что я продолжаю кастинг между списком и ITER. Я также не хочу читать индексы списка, потому что он чувствует себя неэффективно. Может ли кто-нибудь указать мне, как правильно сочетать и сопоставлять itertools?

from itertools import takewhile 

def group(l, p): 
    blocks = [] 
    while True: 
     i = iter(l) 
     taken = list(takewhile(p, i)) 
     l = list(i) 
     if len(taken) > 0: 
      blocks.append(taken) 
     if len(l) == 0: 
      return blocks 

Спасибо!

Answer 1

Вы можете использовать itertools.groupby:

In [22]: import itertools as IT 

In [23]: [list(g) for k, g in IT.groupby(
      [1, 2, 3, 0, 4, 5, 0, 0, 6], lambda x: x != 0) if k] 
Out[23]: [[1, 2, 3], [4, 5], [6]] 

Answer 2

Использование itertools.groupby:

>>> import itertools 
>>> 
>>> data = [1, 2, 3, 0, 4, 5, 0, 0, 6] 
>>> xs = [list(grp) for k, grp in itertools.groupby([1, 2, 3, 0, 4, 5, 0, 0, 6], lambda x: x == 0)] 
>>> xs 
[[1, 2, 3], [0], [4, 5], [0, 0], [6]] 

>>> xs[data[0] == 0::2] 
[[1, 2, 3], [4, 5], [6]] 
>>> xs[data[0] != 0::2] 
[[0], [0, 0]]