Home > front end >  How to group numbers between two occurences of a number using itertools.groupby
How to group numbers between two occurences of a number using itertools.groupby

Time:05-28

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?

CodePudding user response:

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]
  • Related