Search code examples
pythonlistpython-itertools

How to group numbers between two occurences of a number using itertools.groupby


I have a list that looks like this -

nums = [0,0,0,0,1,1,2,3,4,5,6,0,0,0,0,1,2,3,4,5,6,0,0,0,0]

I want to get the numbers between the 0s in the list. For this, I used the code below -

groups = list(itertools.groupby(nums, lambda item:item != 0))
groups = list(filter(lambda item:item[0], groups))
list(map(lambda item:list(item[-1]), groups))

But I am getting empty lists as the output -

[[], []]

My desired output is -

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

How can I do this using itertools.groupby?


Solution

  • Try:

    from itertools import groupby
    
    nums = [0,0,0,0,1,1,2,3,4,5,6,0,0,0,0,1,2,3,4,5,6,0,0,0,0]
    
    output = [list(g) for k, g in groupby(nums, lambda x: x != 0) if k]
    print(output) # [[1, 1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6]]
    

    The doc puts

    The returned group is itself an iterator that shares the underlying iterable with groupby(). Because the source is shared, when the groupby() object is advanced, the previous group is no longer visible.

    Therefore my guest is that the first line

    groups = list(itertools.groupby(nums, lambda item:item != 0))
    

    is the culprit. As list iterates over groupby object, each group in it also gets consumed. You can instead use

    groups = [(k, list(g)) for k, g in groupby(nums, lambda x: x != 0)]
    

    to store the result.


    To see this:

    groups = groupby(nums, lambda x: x != 0)
    k1, g1 = next(groups)
    print(k1, list(g1)) # False [0, 0, 0, 0]
    
    groups = groupby(nums, lambda x: x != 0)
    k1, g1 = next(groups)
    k2, g2 = next(groups)
    print(k1, list(g1), k2, list(g2)) # False [] True [1, 1, 2, 3, 4, 5, 6]